Sync LeetCode submission Runtime - 7 ms (91.36%), Memory - 20.2 MB (33.88%)

Author: jimit105

0023-merge-k-sorted-lists/solution.py

@@ -1,24 +1,40 @@
-import heapq
+# Approach: Heap
+
+# n = total number of nodes, k = no. of linked lists
+# Time: O(n * log k)
+# Space: O(k)
 
 # Definition for singly-linked list.
 # class ListNode:
 #     def __init__(self, val=0, next=None):
 #         self.val = val
 #         self.next = next
 
+import heapq
+
 class Solution:
     def mergeKLists(self, lists: List[Optional[ListNode]]) -> Optional[ListNode]:
-        q = [(l.val, idx) for idx, l in enumerate(lists) if l]
-        heapq.heapify(q)
-        
-        head = curr = ListNode(None)
-        
-        while q:
-            val, idx = heapq.heappop(q)
-            curr.next = ListNode(val)
-            curr = curr.next
-            node = lists[idx] = lists[idx].next
-            if node:
-                heapq.heappush(q, (node.val, idx))
-            
-        return head.next
+        heap = []
+
+        # Add the first node from each list to the heap
+        # We store (value, list_index, node) in heap
+        for i, lst in enumerate(lists):
+            if lst:
+                heapq.heappush(heap, (lst.val, i, lst))
+
+        dummy = ListNode(0)
+        current = dummy
+
+        while heap:
+            val, i, node = heapq.heappop(heap)
+
+            # Add the node to our result list
+            current.next = node
+            current = current.next
+
+            # If there are more nodes in the same list, add the next node to heap
+            if node.next:
+                heapq.heappush(heap, (node.next.val, i, node.next))
+
+        return dummy.next
+