Merging K Sorted Lists

Keep the current head of every list in a min heap. Repeatedly take the smallest head, append it to the result, and push the next element of that list onto the heap. Because the heap always holds the smallest remaining candidate from every list, the next element of the result is always at its root.

import heapq


def merge_sorted(lists):
    heap = []
    positions = [0] * len(lists)
    for index, items in enumerate(lists):
        if items:
            heap.append((items[0], index))
    heapq.heapify(heap)
    result = []
    while heap:
        value, index = heapq.heappop(heap)
        result.append(value)
        positions[index] += 1
        if positions[index] < len(lists[index]):
            heapq.heappush(heap, (lists[index][positions[index]], index))
    return result


print(merge_sorted([[1, 4, 7], [2, 5, 8], [3, 6, 9]]))

The output is 1 through 9 in order. Every element of every list appears exactly once and the result is sorted.

The mechanism is that only the head of each list can be the next element of the merged result, since the rest of that list is at least as large as its head. The heap holds exactly those candidates, so pulling its root gives the next element directly. Each time a list is advanced, its new head takes its place in the heap.

The trade-off to name is against merging the lists pairwise. Merging two at a time costs O of n k in the worst case, because later merges reprocess elements that earlier merges already handled. The heap version is O of n log k, where n is the total number of elements, and the heap never holds more than k entries so the memory is small. If k is small, for example two or three lists, the pairwise version is also fine and simpler to write.

The heapq module provides a merge function that takes any number of sorted iterables and yields their elements in sorted order, using a heap internally.

import heapq


def merge_two(a, b):
    return list(heapq.merge(a, b))


print(merge_two([1, 4, 7], [2, 5, 8]))
print(list(heapq.merge([1, 4, 7], [2, 5, 8], [3, 6, 9])))

The outputs are 1, 2, 4, 5, 7, 8 and 1 through 9. The helper handles any number of inputs, not just two.

The mechanism is the same heap of current heads, but the helper works with iterables rather than lists, so it never needs to know the length of anything. It pulls the smallest candidate, yields it, and advances that iterable. Because the inputs are sorted, the candidate at the root is always the smallest remaining value across all of them.

The trade-off to name is that the helper yields lazily, so it does not build the merged list unless you ask it to. That is exactly what you want for a very large merge, where materialising the result would double the memory. The trade-off in the other direction is that the inputs must already be sorted; if they are not, the output is not sorted and the helper will not warn you. It also assumes comparable items, so mixing types that cannot be compared raises an error.