Top-k and k-way merge

Two common jobs need only a small heap, however big the input: picking the k largest values, and merging k sorted lists. In both, the heap holds k items, so each step costs O(log k) instead of O(log n).

The k largest

To find the k largest values, it seems natural to use a max-heap. Do the opposite: keep the k largest values seen so far in a min-heap. Its root is the smallest of them, the value a newcomer has to beat.

For each new value: if it is not bigger than the root, at least k values seen so far beat it, so it cannot be in the answer; skip it after one comparison. If it is bigger, the root can no longer be in the answer: replace the root with the new value and sift down. Step through the first variant.

import heapq

def top_k(nums, k):
    h = nums[:k]
    heapq.heapify(h)
    for x in nums[k:]:
        if x > h[0]:
            heapq.heapreplace(h, x)  # pop the smallest, push x
    return sorted(h, reverse=True)

The cost is O(n log k) time and O(k) extra space. The input can be a stream too big to hold in memory, such as the lines of a huge log file: only k values are kept. On input in random order the newcomers rarely win, so the work is close to one comparison per value.

Choosing a method

When k is small, the size-k heap is the usual choice. When k is close to n, just sort. Quickselect (DVCQ) is the fastest on average but needs the whole input in memory.

Merging k sorted lists

SORT-2 merged two sorted lists by comparing their front values. With k lists, finding the smallest front value by looking at all k costs O(k) per output value, O(nk) in total. A min-heap of the front values finds it in O(log k):

1. Put the front value of each non-empty list in a min-heap, together with which list it came from.
2. Pop the smallest entry and append its value to the output.
3. If that list has a next value, push it. Repeat until the heap is empty.

Step through the second variant. Each of the n values goes into the heap once and out once: O(n log k) time, and O(k) space besides the output.

def merge_k(lists):
    h = [(l[0], i, 0) for i, l in enumerate(lists) if l]
    heapq.heapify(h)
    out = []
    while h:
        v, i, j = h[0]
        out.append(v)
        if j + 1 < len(lists[i]):
            heapq.heapreplace(h, (lists[i][j + 1], i, j + 1))
        else:
            heapq.heappop(h)
    return out

The entries are (value, list number, position). When two values are equal, the list number settles the order, so the comparison never reaches anything that cannot be compared.

In Python

• heapq.nlargest(k, xs) and heapq.nsmallest(k, xs) use the size-k heap, and also accept a key function, as sorted does. For k = 1 use max or min; when k is near n, use sorted.
• heapq.merge(*iterables) merges lazily: it returns an iterator and reads each input only as far as it needs. That makes it the tool for external sorting: sort chunks of a file too big for memory, write each sorted chunk to disk, then merge the chunks with one pass.