L26 · Heaps & Top-K

Series: LeetCode — From Basics to Interview-Ready · Session 26 / 65 · Phase 1 · Core patterns

Why this session exists

Top K or K smallest → heap. That mapping is the whole session, and it is worth burning in until it is reflex, because the alternative most people reach for — sort everything, then slice the first K — is correct but wasteful, and interviewers notice.

A heap gives you the smallest (or largest) element in O(log n) per operation without ordering everything else. When you only need K of the n elements ranked, sorting all n is doing n log n work to extract a log-k answer. A min-heap of size K does it in O(n log k), and when K is small relative to n — which is the usual case — that is a real win you can state out loud.

The second reason this session exists is that Python's heapq has two sharp edges that trip people in interviews: it is a min-heap only, and it has no built-in max-heap. Both are solved by one trick — negate on the way in, negate on the way out — and once that is muscle memory you stop losing minutes to it.

Blank-file warm-up

Five minutes, empty file, no notes. Reproduce the heapq idioms and the size-K invariant:

import heapq
 
# min-heap (the only kind heapq gives you)
h = []
heapq.heappush(h, x)
smallest = heapq.heappop(h)          # removes and returns the minimum
peek = h[0]                          # minimum without removing
 
# max-heap: negate going in and coming out
heapq.heappush(h, -x)
largest = -heapq.heappop(h)
 
# heapify an existing list in O(n)
heapq.heapify(nums)
 
# K LARGEST -> keep a MIN-heap of size k.
# The root is the smallest of your current top-k; if a new value beats it, swap.
def k_largest(nums, k):
    h = []
    for x in nums:
        heapq.heappush(h, x)
        if len(h) > k:
            heapq.heappop(h)         # evict the smallest -> h holds the k largest
    return h                          # h[0] is the k-th largest

If you cannot type the size-K eviction loop without thinking, that is today's finding. The counter-intuitive part — that you keep a min-heap to find the largest — is exactly the part to over-learn.

Pattern anatomy

The shape: you have n items, you want the K best of them by some key, and K is meaningfully smaller than n.

The invariant: the heap holds the best K seen so far, and its root is the weakest of that current best set. Every new item is compared against the root. If it beats the root, evict the root and insert the new item; the heap size stays at K. If it does not beat the root, it can never be in the top K, so discard it immediately.

The direction flip is the thing to internalise: to keep the K largest, use a min-heap, because you want fast access to the weakest survivor so you can evict it. To keep the K smallest, use a max-heap (negate).

# generic top-k by a key, min-heap of (key, item)
def top_k(items, k, key):
    h = []
    for it in items:
        heapq.heappush(h, (key(it), it))
        if len(h) > k:
            heapq.heappop(h)
    return [it for _, it in h]

The cue

You should hear "heap" from these tells in the problem statement:

1. The literal words "K largest / K smallest / K closest / K most frequent." The most direct cue in all of algorithms. Almost always a heap unless K equals n.
2. "Top K" or "the K best by some score." Same pattern with a key function.
3. A stream you cannot fully sort — elements arrive one at a time and you must maintain a running top-K without storing and re-sorting everything.
4. "Merge K sorted lists/arrays" — a min-heap of the K current heads pulls the global minimum in O(log k) each step.
5. "Median of a stream" — the two-heap variant (you will meet this again in L52). If you see running median, think heaps, not sorting.

Guided solve

Kth Largest Element in an Array. Return the k-th largest element (in sorted order, not distinct).

Narrate it: the k-th largest is exactly the smallest element of the top-K set. So keep a min-heap of size K; after processing every element, its root is the answer.

def findKthLargest(nums, k):
    h = []
    for x in nums:
        heapq.heappush(h, x)
        if len(h) > k:
            heapq.heappop(h)
    return h[0]          # smallest of the k largest == k-th largest overall

O(n log k) time, O(k) space. The instinct to sort and index nums[-k] also works at O(n log n) — say so, then explain why the heap is better when k is small. (A quickselect solution reaches O(n) average and is worth mentioning as the "if they push on complexity" answer.)

Solo timed

Fifteen minutes each, timer visible. Hints only — no solutions.

• Top K Frequent Elements — count with a Counter first, then run the size-K heap on (count, value) pairs. heapq.nlargest(k, counter, key=counter.get) is the one-liner, but implement the heap yourself at least once.
• Merge k Sorted Lists — push the head of each list as (value, list_index, node) (the index breaks ties so nodes never get compared). Pop the min, push its successor. O(N log k).
• K Pairs with Smallest Sums — a min-heap seeded with the first column; each pop lazily pushes its right neighbour. Do not generate all pairs.

Common failure modes

1. Using a max-heap for K-largest. The classic inversion error. K-largest wants a min-heap so the root is the evictable weakest. Wire it backwards and you evict your best elements.
2. Forgetting Python has no max-heap. You negate, or you spend the interview confused about why heappop gives the wrong end.
3. Comparing un-comparable payloads. Pushing (count, node) where two counts tie makes Python compare the nodes and crash. Insert a unique tiebreaker index: (count, i, node).
4. Sorting when a heap was the point. sorted(nums)[-k] passes the tests but signals you did not see the pattern. It is O(n log n), not O(n log k).

Complexity

The size-K heap is O(n log k) time — n insertions, each O(log k) because the heap never exceeds K — and O(k) space. Compare against sort-and-slice at O(n log n) time and O(n) space. For merge-K-lists the bound is O(N log k) where N is the total node count and k the number of lists.

Spaced queue

Re-solve whatever is due from earlier sessions before you close. Status ladder for everything you touched today:

• Cold — unaided, under 15 minutes, optimal → returns in 60 days
• Warm — unaided but slow or suboptimal → returns in 21 days
• Hint — needed a nudge or a peek → returns in 7 days
• Failed — read the editorial → returns in 2 days, then 7

Kth Largest and Top K Frequent are the two to over-learn; they are the base cases the rest of the family reduces to.