Merge k Sorted Lists

Problem Id: 23 Difficulty: Hard Tag: Linked List Tag: Divide and Conquer Tag: Heap


Intuition

There are k linked lists, so we will need k indexes. Each time, if we simply loop on k indexes to find the smallest node, the overall time complexity would be O(n * k). To reduce the time complexity, we can use a heap to store current visiting node for each linked list, pop the smallest one and push it's next node each time. So the overall time complexity would be O(n * lg(k)).

Solution


# 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[ListNode]) -> ListNode:
        heap = []
        for group, head in enumerate(lists):
            if head is not None:
                heapq.heappush(heap, (head.val, group))
        currents = lists[:]

        head = ListNode(None)
        current = head
        while heap:
            _, group = heapq.heappop(heap)
            node = currents[group]
            if node.next:
                heapq.heappush(heap, (node.next.val, group))
                currents[group] = node.next
            node.next = None
            current.next = node
            current = node

        return head.next