Merge k Sorted Lists https://leetcode.com/problems/merge-k-sorted-lists/
Use min heap of size k to merge: O(Nlog(k))
- Maintain a min heap of size k and add the heads of each list to the heap.
- Now pop the heap and add to result. Then add to heap the next node for the popped element.
from heapq import heappush, heappop class Solution(object): def add_to_heap(self, i, heap, lists): if lists[i]: temp, lists[i].next = lists[i].next, None heappush(heap, (lists[i].val, (i, lists[i]))) lists[i] = temp return def mergeKLists(self, lists): """ :type lists: List[ListNode] :rtype: ListNode """ N, heap, result = len(lists), , ListNode(-1) curr = result for i in range(N): self.add_to_heap(i, heap, lists) while len(heap): x, (i, lnode) = heappop(heap) curr.next, curr = lnode, lnode # Statement makes sure old values are assigned. self.add_to_heap(i, heap, lists) return result.next