Heaps and top-k

A binary heap is a tree stored in an array, with one rule: each parent is smaller than its children (a min-heap) or larger (a max-heap). The root is always the smallest or largest item. Pushing an item and popping the root each take O(log n), and reading the root takes O(1).

Heaps matter most for a job that comes up constantly: "give me the top k of a large stream," such as the 10 slowest requests, the 5 biggest orders or the next task due.

The big idea

Sorting a million items to find the top ten costs about n log n. A heap of size k finds them in about n log k, which for k = 10 is barely more than a single scan. And it uses O(k) memory, so it works on a stream you can't store.

The trick is to keep a min-heap of size k. Its root is the smallest of the top k so far. Each new item either beats that root, and then replaces it, or it doesn't, and then it's ignored.

Walking through it

import heapq

def top_k(items, k):
    heap = []                                    # min-heap of the k largest seen so far
    for x in items:
        heapq.heappush(heap, x)
        if len(heap) > k:
            heapq.heappop(heap)                  # drop the smallest
    return sorted(heap, reverse=True)

print(top_k(range(1, 11), 3))

Output:

[10, 9, 8]

Each item is pushed once. The heap never grows beyond k + 1 items, so the work per item is O(log k). When the heap has more than k items, the smallest is removed. At the end, the heap holds exactly the k largest values, and sorting that small heap gives them in order.

For a stream, the same loop runs without storing the input. For a leaderboard that is updated over time, keep the heap and push each new score.

Heaps in systems

Heaps show up in schedulers that run the next timer or the next job by due time, in Dijkstra's shortest-path algorithm for the closest unvisited node, and in merging sorted streams. Timers in event loops, covered in event loops, are usually a heap or a timing wheel, so that the next deadline is always at the root.

Common mistakes

• Sorting the whole input to get the top k, when the input is large or streamed.
• Using a max-heap for top-k, which holds the wrong items. Use a min-heap of size k, so the root is the cutoff.
• Forgetting that a heap doesn't support finding an arbitrary item quickly. For that, use an indexed structure or a tree.

Check yourself

• Why does a min-heap of size k give the top k, and not the bottom k? Answer: its root is the smallest of the k kept items, so any larger new item can replace it. The kept items are the largest seen.
• Why is top-k with a heap cheaper than sorting? Answer: each item costs O(log k) against the heap, instead of O(log n) for a full sort, and the memory stays at k items.

Key takeaways

• A binary heap keeps the smallest or largest item at the root in O(1), with O(log n) updates.
• Top-k uses a min-heap of size k, so it finds the largest k in one pass.
• The memory stays at k items, which suits streams.

Further learning

Next, see sorted trees, skip lists and tries for structures that keep keys ordered for range queries and prefix searches.