Merge k Sorted Lists

HardHeapLeetCode 23 ↗World 9-10
0:00 / 0:00

Merge k sorted linked lists into one. The next node is always one of the k heads, so keep them in a min-heap and pop the smallest.

▼

The problem

LeetCode 23 (Hard). Given an array of k linked lists, each sorted in ascending order, merge them into one sorted linked list and return its head.

Example: lists = [[1,4,5],[1,3,4],[2,6]] → [1,1,2,3,4,4,5,6].

TRY IT ON LEETCODE ▶

The solution

import heapq

def mergeKLists(lists):
    heap = []
    for i, node in enumerate(lists):
        if node:
            heapq.heappush(heap, (node.val, i, node))
    dummy = tail = ListNode()
    while heap:
        val, i, node = heapq.heappop(heap)
        tail.next = node
        tail = node
        if node.next:
            heapq.heappush(heap, (node.next.val, i, node.next))
    return dummy.next

Transcript

Merge k Sorted Lists. You get k linked lists, each one sorted. Merge them into one sorted list, and return its head.

Take three lists: one, four, five. One, three, four. And two, six. Merged: one, one, two, three, four, four, five, six.

The naive way merges the lists one at a time into a growing result. With N nodes in all, the early nodes get walked again on every merge: order k times N. Sorting all the values is order N log N, and ignores that each list is sorted.

Better: the next node is always one of the k heads. So keep the heads in a min-heap. Pop the smallest, append it, and push that node's next. The heap never holds more than k nodes.

In code, push each head as value, list index, node. Equal values compare by index, so two nodes are never compared. While the heap has nodes: pop, append, and push the next node if there is one.

On the example, heads one, one and two go in. Pop one from list zero, push four. Pop one from list one, push three. Pop two, push six. Pop three, push four. The fours tie, so list zero goes first and pushes five. Then four, five, six. Done.

Each of the N nodes is pushed and popped once, at log k each. So time is order N log k, and space is order k for the heap. Merging the lists in pairs, round by round, also gets N log k.

Keep every head in a heap, and always take the smallest. That's Merge k Sorted Lists.