Merge k sorted lists: min-heap over the current heads

Problem. Given k sorted linked lists, merge them into one sorted list and return its head.

This is the classic heap-based multiway merge pattern.

Examples:

  • lists = [[1,4,5],[1,3,4],[2,6]][1,1,2,3,4,4,5,6]
  • lists = [][]
  • lists = [[],[1]][1]

Solution. Keep the current head of each list in a min-heap. Repeatedly extract the smallest node, append it to the answer, and push that node’s successor if it exists.

import heapq


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


def merge_k_lists(lists):
    heap = []
    counter = 0

    for node in lists:
        if node:
            heapq.heappush(heap, (node.val, counter, node))
            counter += 1

    dummy = ListNode()
    tail = dummy

    while heap:
        _, _, node = heapq.heappop(heap)
        tail.next = node
        tail = tail.next
        if node.next:
            heapq.heappush(heap, (node.next.val, counter, node.next))
            counter += 1

    return dummy.next

O(n log k) time for n total nodes and k lists, O(k) heap space.

Gotchas.

  • In Python, add a counter tie-breaker because ListNode objects are not directly orderable.
  • Pairwise merging also works but is slower if done naively from left to right.
  • The heap stores at most one node from each list at a time.

Run: Merge K Sorted Lists