Top-k and two heaps

Three problems account for almost every heap that appears in an interview: keep the best k, merge m sorted streams, and track a statistic that lives in the middle of a growing set. All three work because a heap gives you the extreme element in O(1) and a repair in O(log n), and all three are compared against the same alternative — just sort it.

Top-k: a min-heap of size k

For the k largest values, hold a min-heap of size k. The inversion looks wrong for about two seconds and is the whole trick: the top of that heap is the weakest survivor, so one comparison decides whether a new value belongs.

import heapq

def top_k(stream, k):
    h = []
    for x in stream:
        if len(h) < k:
            heapq.heappush(h, x)
        elif x > h[0]:                 # beats the weakest survivor
            heapq.heapreplace(h, x)    # pop the weakest, push x, one sift
    return sorted(h, reverse=True)

Cost: n comparisons against h[0], plus at most n sifts of log₂ k each — O(n log k). Sorting is O(n log n). At n = 1,000,000 and k = 10:

sort:      1,000,000 × log₂(10⁶) = 1,000,000 × 20  ≈ 20,000,000 comparisons
size-k heap: 1,000,000 × log₂(10)  ≈ 1,000,000 × 3.3 ≈  3,300,000

Six times fewer, and that ratio is the least interesting part. The heap holds 10 values — about 80 bytes — while the sort holds 8 MB, and most values never sift at all because they fail x > h[0] in one comparison. The heap is what lets you do this on a stream you never store: a log tailer or a Kafka consumer can track the top 10 slowest requests over a billion events in constant memory. Sorting cannot start until the last event arrives.

Two boundaries worth naming. When k approaches n, log k approaches log n and the heap wins nothing — sort. And when the data is already in memory and you want the k largest as an unordered set, quickselect is O(n) average, beating O(n log k); the heap wins when the data is a stream, when you want the result maintained as data keeps arriving, or when you cannot afford to mutate the input.

Python ships this as heapq.nlargest(k, data), which does exactly the above.

K-way merge: one entry per stream

Given m sorted sequences totalling n items, produce one sorted sequence. The naive version concatenates and sorts — O(n log n), and it needs all n items in memory at once, which fails the moment the inputs are files larger than RAM.

Instead, hold a heap of exactly m entries: the current head of each stream. Pop the smallest, emit it, and push that stream's next value.

import heapq

def merge_k(lists):
    h = [(lst[0], i, 0) for i, lst in enumerate(lists) if lst]
    heapq.heapify(h)                      # O(m), not O(m log m)
    out = []
    while h:
        val, i, j = heapq.heappop(h)
        out.append(val)
        if j + 1 < len(lists[i]):
            heapq.heappush(h, (lists[i][j + 1], i, j + 1))
    return out

Each item is pushed once and popped once, so the cost is O(n log m). With m = 100 sorted files and n = 1,000,000 records:

concatenate + sort: 1,000,000 × log₂(10⁶) = 20,000,000 comparisons, 1,000,000 records resident
100-entry heap:     1,000,000 × log₂(100) ≈  6,600,000 comparisons,       100 records resident

Three times fewer comparisons, and 10,000× less memory. That second column is why external sorting works at all: sort chunks that fit in RAM, write them out, then merge them with a heap whose entire working set is one record per file. Note the tuple carries the stream index i as a tie-breaker, so equal values never make Python compare the payloads — the TypeError trap from the heap invariant.

heapq.merge is the library version, and it returns an iterator, so the output never has to be resident either.

Two heaps: a running median

The median is not an extreme, so one heap cannot hold it. Two can, if you split the data at the middle and point the heaps at each other: a max-heap for the lower half, a min-heap for the upper half. The median is then one or two peeks at the two tops.

A max-heap of the lower half and a min-heap of the upper half, with every new value pushed low, spilled high, then rebalanced back100%−+FitSVGEsc

Scroll to zoom · drag to pan · 0 fits · Esc closes

The rebalance is where implementations go wrong, so make it unconditional rather than clever. Every value takes the same route: push into the low heap, move the low heap's top into the high heap, then move it back only if the high heap has become the longer one.

import heapq

class RunningMedian:
    def __init__(self):
        self.lo = []          # max-heap of the lower half, values negated
        self.hi = []          # min-heap of the upper half

    def add(self, x):
        heapq.heappush(self.lo, -x)                      # 1 always push low
        heapq.heappush(self.hi, -heapq.heappop(self.lo)) # 2 spill the top up
        if len(self.hi) > len(self.lo):                  # 3 keep lo ≥ hi
            heapq.heappush(self.lo, -heapq.heappop(self.hi))

    def median(self):
        if len(self.lo) > len(self.hi):
            return -self.lo[0]
        return (-self.lo[0] + self.hi[0]) / 2

Two invariants do the work, and you should state both out loud. Order: every value in lo is ≤ every value in hi, which step 2 guarantees — the value that crosses is always the largest of the lower half. Size: len(lo) equals len(hi) or exceeds it by exactly one, which step 3 restores. Together they put the median at lo[0] when the count is odd and between the two tops when it is even.

Three moves per value — push, spill, and a rebalance on every other value — which is three to five heapq calls, averaging four: step 3 fires on exactly the inserts that leave hi longer, which is every second one. Each call is O(log n), and the query is O(1). Against the alternatives, over a stream of 100,000 values with the median read after each one:

re-sort each time:        100,000 × 100,000 × 17    ≈ 1.7 × 10¹¹ comparisons
insert into sorted list:  100,000 × 50,000          = 5 × 10⁹ element moves
two heaps:                100,000 × 4 × log₂(10⁵)   ≈ 6.6 × 10⁶ operations

The sorted-list row is the one that catches people, because the binary search for the insertion point really is O(log n) — but the insert itself shifts half the array, and 5 × 10⁹ moves at roughly 1 ns each is about five seconds of work that the heaps do in milliseconds.

The same two-heap frame answers any split statistic: for the 90th percentile, keep the size ratio at 9:1 instead of 1:1 and read hi[0]. What it does not handle is removal from the middle — a sliding-window median needs lazy deletion with a map of tombstones, or an order-statistic tree.

In an interview

State the heap's size before you write a line. "A min-heap of size k, so the top is the smallest of the k largest, and I compare each incoming value against it" is the sentence being graded; the code after it is bookkeeping. Then give the cost as O(n log k) and name why it is not O(n log n).

For the median, say the two invariants — order and size — and then say which step of the code maintains each. An interviewer who hears "step 2 guarantees the order, step 3 the balance" does not need to read the rest.

The mistake that loses points: using a max-heap for the k largest. It gives you the biggest value in O(1), which is not the question — you need the smallest of the k you kept, so you can evict it. Candidates who reach for the max-heap end up holding all n elements and have quietly rebuilt sorting.

The near-miss on the median is a conditional rebalance: "if x is less than the top of lo, push it there, else push it to hi, then fix the sizes." That is correct if you get all four branches right, and most people get three. The unconditional push-spill-rebalance has one path and costs about two extra heap operations per value.

Check yourself

You need the 100 most-viewed items from a stream of 1 billion events. What do you keep, and what does bounding the memory cost you?

Not a 100-entry heap on its own. An event carries an item, not that item's total, so there is nothing to compare against h[0] yet; and an item evicted from the heap takes its count with it, so when it reappears you cannot tell it from a new one. Exact top-k by frequency needs a count per distinct item: a hash map of item → count over the billion events, then one pass of that map into a size-100 min-heap. Memory is proportional to distinct items, not to k — 10 million distinct items at tens of bytes an entry is hundreds of megabytes. Bounding it below that costs exactness: Space-Saving keeps a fixed 1,000 counters and overwrites the weakest with the new item, inheriting its count, and a count-min sketch estimates counts in fixed memory. Both can over-report, so both give an approximate top-100. The size-k heap in the body applies when the ranking key arrives with the event — the 100 slowest of a billion latency samples — where one comparison against h[0] really does decide.

Your median class pushes to hi when x is greater than hi[0] and to lo otherwise, then rebalances sizes. Is it correct? Is it better?

It is correct if every branch is right, and it does roughly halve the heap operations, about four per value down to two: one call when no rebalance is needed, three when one is. It is worse to write under time pressure, because it has four paths to reason about, including the empty-heap case where hi[0] does not exist. The unconditional version has one path.

You are merging 500 sorted files, 2 GB each, on a box with 8 GB of RAM. What is in memory, and what is the cost?

A 500-entry heap of one record per file, plus a read buffer per file — kilobytes, not gigabytes. The cost is O(n log 500) ≈ 9 comparisons per record, against log₂ n ≈ 30 or more if you could sort the whole terabyte at once, which you cannot. Merging is the only version that fits, which is the point made in complexity by counting: the constraint picks the algorithm before the complexity does.