Refactored Merge K Sorted Lists.java

Author: varunu28

Hard/Merge K Sorted Lists.java

@@ -3,41 +3,29 @@
  * public class ListNode {
  *     int val;
  *     ListNode next;
- *     ListNode(int x) { val = x; }
+ *     ListNode() {}
+ *     ListNode(int val) { this.val = val; }
+ *     ListNode(int val, ListNode next) { this.val = val; this.next = next; }
  * }
  */
 class Solution {
-    public ListNode mergeKLists(ListNode[] lists) {
-        if (lists.length == 0) {
-            return null;
-        }
-
-        PriorityQueue<ListNode> queue = new PriorityQueue<>(lists.length, new Comparator<ListNode>() {
-            @Override
-            public int compare(ListNode o1, ListNode o2) {
-                return o1.val - o2.val;
-            }
-        });
-
-        for (int i = 0; i < lists.length; i++) {
-            if (lists[i] != null) {
-                queue.add(lists[i]);
-            }
-        }
-
-        ListNode head = new ListNode(0);
-        ListNode p = head;
-
-        while (!queue.isEmpty()) {
-            ListNode node = queue.poll();
-            p.next = node;
-            if (node.next != null) {
-                queue.add(node.next);
-            }
-
-            p = p.next;
-        }
-
-        return head.next;
+  public ListNode mergeKLists(ListNode[] lists) {
+    PriorityQueue<ListNode> pq = new PriorityQueue<>((a, b) -> a.val - b.val);
+    for (ListNode listNode : lists) {
+      if (listNode != null) {
+        pq.add(listNode);
+      }
     }
+    ListNode dummy = new ListNode(0);
+    ListNode curr = dummy;
+    while (!pq.isEmpty()) {
+      ListNode removed = pq.remove();
+      curr.next = new ListNode(removed.val);
+      curr = curr.next;
+      if (removed.next != null) {
+        pq.add(removed.next);
+      }
+    }
+    return dummy.next;
+  }
 }