Heaps and Top-K

TL;DR: A heap is a priority queue: peek the min (or max) in O(1), insert and pop in O(log n). For "top K of a stream of n", keep a min-heap of size K; push each item, and if the heap grows past K, pop the smallest. You end with the K largest in O(n log k) time and O(k) space, without sorting the whole input.

Why a heap instead of sorting

If you need the full ranking, sort: O(n log n). But most "top-K" questions only need the K best, and K is tiny next to n (the 10 worst latencies out of 10 million requests). A full sort wastes work ordering the 9,999,990 you will throw away. A bounded heap pays O(log k) per element and keeps only K in memory, which is what lets it run over a stream that does not fit in RAM.

The two approaches to top-K, side by side:

The move that trips people up: to find the K largest, you keep a min-heap. The root is then the smallest of your current top-K, so it is exactly the element to evict when a bigger one arrives.

Top-K from a stream

import heapq

def top_k(stream, k: int) -> list[int]:
    heap: list[int] = []          # min-heap, size <= k
    for x in stream:
        if len(heap) < k:
            heapq.heappush(heap, x)
        elif x > heap[0]:         # bigger than the current smallest kept
            heapq.heapreplace(heap, x)   # pop min, push x, in one O(log k)
    return sorted(heap, reverse=True)

assert top_k([5, 1, 9, 3, 7, 2, 8], 3) == [9, 8, 7]
assert top_k([4, 4, 4], 2) == [4, 4]
assert top_k([], 3) == []

Each element costs at most O(log k), so the scan is O(n log k) with O(k) memory. Python's heapq is a min-heap; for a max-heap, push negated keys or use tuples. heapreplace is the right primitive here: it pops then pushes in a single sift, which is why the x > heap[0] guard comes first.

Merging K sorted sources

The same structure merges K sorted lists or streams in O(n log k): seed the heap with the head of each source, pop the smallest, and push the next item from whichever source it came from.

import heapq

def merge_sorted(lists: list[list[int]]) -> list[int]:
    heap = [(lst[0], i, 0) for i, lst in enumerate(lists) if lst]
    heapq.heapify(heap)           # O(k)
    out = []
    while heap:
        val, li, idx = heapq.heappop(heap)
        out.append(val)
        if idx + 1 < len(lists[li]):
            heapq.heappush(heap, (lists[li][idx + 1], li, idx + 1))
    return out

assert merge_sorted([[1, 4, 7], [2, 5], [3, 6, 9]]) == [1, 2, 3, 4, 5, 6, 7, 9]

The tuple carries (value, list_index, element_index); the extra fields both break ties deterministically and tell you where to pull the next element.

Common pitfalls

• Max-heap confusion. To keep the K largest, use a min-heap of size K and evict its root. Reaching for a max-heap of the whole input defeats the memory win.
• Comparing raw objects. Heaps compare by the first tuple element; if values tie, include a unique tiebreaker (an index) so Python never tries to compare unorderable payloads.
• nlargest for large K. heapq.nlargest(k, data) is perfect for one-shot top-K, but for k close to n a plain sorted() is simpler and not slower.
• Forgetting the heapify shortcut. Building a heap from a list is O(n) via heapify, cheaper than n individual pushes.

Key takeaways

• Heap operations: peek O(1), push/pop O(log n). Use it whenever "the best so far" must stay cheap to query.
• Top-K of n is O(n log k) with a size-K heap, beating an O(n log n) sort when k is small or the data streams.
• K largest needs a min-heap (evict the root); K smallest needs a max-heap. Say which and why.
• Merging K sorted sources is the same pattern at O(n log k); carry source indices in the heap tuple for tie-breaking and bookkeeping.