Merge k Sorted Lists

You'll keep the front node of every list in a min-heap, so the smallest is always one pop away. You need linked list nodes and Python's heapq. Related chapter: heaps.

1. The problem

You get k sorted linked lists. Merge them into one sorted linked list, keeping every value, including repeats. Some lists may be empty. Our version reuses the input nodes and changes their next links. This is LeetCode 23. The tables below show each list as its values.

2. A correct baseline

class ListNode:
    def __init__(self,val=0,next=None):
        self.val,self.next=val,next

def build(values):
    head=None
    for v in reversed(values): head=ListNode(v,head)
    return head

def values(head):
    result=[]
    while head:
        result.append(head.val); head=head.next
    return result

def baseline(lists):
    return build(sorted(v for head in lists for v in values(head)))

assert values(baseline([build([1,4]),build([1,3]),build([2])]))==[1,1,2,3,4]
assert values(baseline([None,build([5])]))==[5]
assert baseline([]) is None

Pulling out every value and sorting takes O(N log N) time for N nodes in total, plus O(N) extra space. On a million nodes, that's about 20 million comparisons.

The bottleneck. It sorts everything from scratch, ignoring that each list is already sorted.

3. The observation

Take [1,4] and [2,3]. The smallest value overall must be at the front of some list: here, 1 or 2. Take the 1. Now the fronts are 4 and 2, and the next smallest is again one of the fronts.

So you only ever need to compare the front node of each list. Keep those fronts in a min-heap, which hands back the smallest in O(log k) time. Pop the smallest, add it to the output, and push the next node from its list. Each heap entry also stores its list's index. When two values tie, Python compares the index next, instead of trying to compare two nodes, which would raise an error.

4. The solution and a trace

import heapq

class ListNode:
    def __init__(self,val=0,next=None):
        self.val,self.next=val,next

class Solution:
    def mergeKLists(self, lists: list) -> object:
        """Sorted acyclic lists; reuse their nodes in the merged chain."""
        heap=[(node.val,i,node) for i,node in enumerate(lists) if node]
        heapq.heapify(heap)
        dummy=ListNode(); tail=dummy
        while heap:
            _,i,node=heapq.heappop(heap)
            following=node.next
            tail.next=node; tail=node
            if following:
                heapq.heappush(heap,(following.val,i,following))
        tail.next=None
        return dummy.next

solve=Solution().mergeKLists
assert values(solve([build([1,4]),build([1,3]),build([2])]))==[1,1,2,3,4]
assert values(solve([None,build([5])]))==[5]
assert solve([]) is None

Trace [[1,4],[2,3]]:

5. Why it works

What stays true. The heap holds the first unused node of every list that isn't finished.

The output is sorted. Each list is sorted, so the smallest front value is the smallest value left anywhere. Taking it next keeps the output sorted. Pushing its next node keeps the heap holding every list's front.

Nothing is lost or repeated. Each node enters the heap once and leaves once, so each appears in the output exactly once.

6. Cost, edge cases, and failure

Cost. Building the heap takes O(k) time. Each of the N nodes is popped and pushed once, at O(log k) each. That's O(N log k) total. With a million nodes in 10,000 lists, that's about 14 steps per node, instead of about 20 for sorting everything. The heap holds at most k nodes, so O(k) extra space. The only new node is the dummy placeholder at the start.

Edge cases. Empty lists are skipped when the heap is built. LeetCode defines ListNode for you; this lesson defines it so the code runs. The lists must not share nodes or contain cycles.

Where it fails. If an input list isn't sorted, its front isn't its smallest, and the output comes out unsorted:

assert sorted([4,1])==[1,4]
assert values(solve([build([4,1])]))==[4,1]  # Deliberately unsorted input breaks the frontier argument.

import random
rng=random.Random(23)
for _ in range(250):
    arrays=[sorted(rng.randint(-5,5) for _ in range(rng.randrange(8))) for _ in range(rng.randrange(8))]
    expected=values(baseline([build(a) for a in arrays]))
    assert values(solve([build(a) for a in arrays]))==expected