Heap & Top-K

TL;DR

If you need the k largest items out of n, keep a min-heap of size k. Push every incoming value; whenever the heap grows past k, pop the smallest. What remains is the top-k. The root of the heap is the kth largest — the weakest survivor — and it’s the threshold every new candidate must beat. The dual problem (k smallest) flips the heap: a max-heap of size k, popping the largest whenever the size exceeds k.

Python’s heapq is a min-heap only. For a max-heap, push -x and negate on read. For top-k by a custom score, push tuples (score, tiebreaker, payload) — heapq compares lexicographically, so the tiebreaker is what saves you when two scores collide and the payloads aren’t comparable.

The streaming twist: a two-heap median keeps a max-heap of the lower half and a min-heap of the upper half, rebalanced so their sizes differ by at most one. The median is either the top of the larger heap (odd count) or the average of both tops (even count). Every insert is O(log n).

The bound to remember: top-k over n items is O(n log k), not O(n log n). When k ≪ n (k=10 over a billion log lines), that’s the difference between a streaming aggregator and a batch sort.

When to reach for it

Signals that a heap is the right tool:

• The question asks for “top k”, “kth largest/smallest”, “k closest”, or “k most frequent” — and k is much smaller than n.
• The data is a stream: items arrive one at a time, and you need an always-current top-k or running median without rescanning history.
• You need partial order, not full order. Sorting all n items wastes work if you only care about the top slice.
• You’re merging k sorted streams (log files, posting lists, shard results) — a min-heap keyed on the head of each stream is the canonical k-way merge.

It’s not a heap problem if you need the items in sorted order at the end (then sort), if k is within a constant factor of n (sort-then-slice is clearer), or if you need O(1) lookup by key (use a hash map or an indexed priority queue).

Flavor 1 — kth largest in array

Given nums and k, return the kth largest element. Maintain a min-heap of size k. The root is always the kth largest seen so far.

import heapq

def kth_largest(nums: list[int], k: int) -> int:
    heap: list[int] = []
    for x in nums:
        if len(heap) < k:
            heapq.heappush(heap, x)
        else:
            heapq.heappushpop(heap, x)  # push then pop the smallest
    return heap[0]

heappushpop(heap, x) is the workhorse: if x is smaller than the current min, it returns x unchanged; otherwise it pushes x and pops the old min. One log-k operation instead of a push followed by a separate pop.

Worked example. nums = [3, 2, 1, 5, 6, 4], k = 2. We want the 2nd largest, which is 5.

Final heap [5, 6], root is 5. The 2nd largest. Note that step 3 doesn’t change the heap at all — 1 is smaller than the current threshold 2 and gets discarded immediately. That early rejection is the whole reason this beats a full sort.

The same skeleton solves top-k frequent (LC 347): build a Counter, then push (count, item) tuples through the same min-heap of size k.

Flavor 2 — streaming top-k (LC 703)

Same idea, but the stream is open-ended and add returns the current kth largest after each insert. Stash the heap on self.

import heapq

class KthLargest:
    def __init__(self, k: int, nums: list[int]):
        self.k = k
        self.heap = []
        for x in nums:
            self.add(x)

    def add(self, val: int) -> int:
        if len(self.heap) < self.k:
            heapq.heappush(self.heap, val)
        else:
            heapq.heappushpop(self.heap, val)
        return self.heap[0]

Trace. k = 3, init [4, 5, 8, 2], then add(3), add(5), add(10), add(9), add(4).

After init the heap holds the three largest of [4, 5, 8, 2]: [4, 5, 8], root 4. add(3) → pushpop(3) returns 3, heap unchanged, return 4. add(5) → pushpop(5) evicts the old 4, heap becomes [5, 8, 5], return 5. add(10) → evicts 5, heap [5, 8, 10], return 5. add(9) → evicts 5, heap [8, 9, 10], return 8. add(4) → discarded, return 8.

Each add is O(log k). The init is O(n log k) where n is the seed length.

Flavor 3 — two-heap median (LC 295)

Maintain two heaps that split the data at the median:

• lo — max-heap of the lower half (in Python: store negatives).
• hi — min-heap of the upper half.

Invariants: every element of lo is ≤ every element of hi, and len(lo) - len(hi) ∈ {0, 1}. With lo allowed to be one larger, the median is -lo[0] for odd counts and (-lo[0] + hi[0]) / 2 for even.

import heapq

class MedianFinder:
    def __init__(self):
        self.lo: list[int] = []  # max-heap, store negatives
        self.hi: list[int] = []  # min-heap

    def addNum(self, num: int) -> None:
        # route the new value into the correct half
        if self.hi and num > self.hi[0]:
            # belongs in the upper half; push there, migrate the new min back to lo
            heapq.heappush(self.lo, -heapq.heappushpop(self.hi, num))
        else:
            heapq.heappush(self.lo, -num)
        # rebalance: lo can be one bigger than hi, never more
        if len(self.lo) > len(self.hi) + 1:
            heapq.heappush(self.hi, -heapq.heappop(self.lo))
        elif len(self.hi) > len(self.lo):
            heapq.heappush(self.lo, -heapq.heappop(self.hi))

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

Worked example. Stream 1, 2, 3, 4, 5. lo shown as the conceptual max-heap (negate the stored values for the actual array).

Each insert: at most two heap ops, each O(log n), so O(log n) per addNum and O(1) per findMedian.

Why it works — the invariant

A min-heap of size k always holds the largest k values seen so far. The proof is a one-liner: when a new value x arrives, either x is bigger than the smallest survivor (root) — in which case the root is not in the top k anymore and gets evicted — or x is smaller, in which case x itself can’t be in the top k and gets discarded. Either way, the heap’s contents after the operation are exactly the top k of everything seen.

Each heap op costs O(log k) because the tree’s height is log k. Across n insertions, total work is O(n log k). The log k factor — not log n — is what makes this pattern shine when k is small.

Complexity

• Time: O(n log k) to build the top-k from a stream of n. Each add / addNum is O(log k) (or O(log n) for the two-heap median, since both heaps grow without bound).
• Space: O(k) for top-k, O(n) for the median (you keep every element).
• Heapify shortcut: if you have all n items up front and want the top k, heapq.nlargest(k, nums) is O(n log k) and reads cleaner than the manual loop. Use the manual loop only when the stream is online.

Common pitfalls

• heapq is min-only. For a max-heap, push -x (or for objects, wrap in a class with __lt__ reversed). Forgetting the negation flips your answer.
• Tiebreakers in tuples. heapq.heappush(h, (score, item)) blows up when two scores tie and item isn’t orderable (e.g. dicts, custom objects). Always include a tiebreaker: (score, counter, item) where counter is a monotonically increasing int.
• Two-heap drift. Forgetting the rebalance step lets one heap grow arbitrarily larger than the other; the median formula then reads from the wrong root. Rebalance after every insert, not just occasionally.
• Wrong heap direction. “Kth largest” wants a min-heap (so the smallest of the survivors — the threshold — is at the root and cheap to evict). Using a max-heap of size k for kth-largest is a common bug; you’d evict the wrong end every time.
• Sort-then-slice when it’s fine. If k is within a constant factor of n, sorted(nums)[-k:] is O(n log n) but reads in one line. Don’t reach for a heap to save a log factor that doesn’t matter.

Where you see this in production

• Lucene’s top-N collector. Search engines never sort all matching documents — they keep a min-heap of size N (the page size) keyed on relevance score. Documents below the heap’s root are pruned without scoring further. Elasticsearch and Solr inherit this directly from Lucene’s TopScoreDocCollector.
• Kafka partition leader balancing. The controller picks under-replicated partitions to reassign by maintaining a priority queue of brokers keyed on current leader load — the “least loaded broker” pop is a min-heap op on every reassignment decision.
• Redis sorted sets. ZADD / ZRANGEBYSCORE are backed by a skip list plus hash, but the contract — “give me the top-K by score, streaming” — is the heap interface, and Redis is what people reach for when they need streaming top-k over a distributed event stream (leaderboards, rate limiters, recent-N feeds).
• Linux CFS scheduler. The Completely Fair Scheduler picks the next task by popping the leftmost node of a red-black tree keyed on virtual runtime — a balanced-tree analogue of the priority queue, with the same “always pick the minimum” semantics a heap would give you.

Practice problems

Resources

• NeetCode roadmap — neetcode.io/roadmap — pattern-by-pattern problem sets organised exactly like this site.
• NeetCode practice grid — neetcode.io/practice — track which problems you’ve solved per pattern.
• NeetCode YouTube — @NeetCode — clear, whiteboard-style walkthroughs for almost every LeetCode problem above.