L26 · Heaps & Top-K

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

Watch first

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.)

Top K Frequent Elements, end to end. Given nums and k, return the k most frequent values.

Brute force: count, then sort the counts descending, then slice — O(n + m log m) where m is the number of distinct values. Correct, and the interviewer will immediately say "better than O(n log n), please."

The pattern insight: you do not need the distinct values ordered, only the top k of them, so cap a min-heap at k keyed on count.

from collections import Counter
import heapq
 
def topKFrequent(nums, k):
    counts = Counter(nums)              # O(n), m distinct keys
    h = []                              # min-heap of (count, value), capped at k
    for value, c in counts.items():
        heapq.heappush(h, (c, value))   # tuples compare on count first
        if len(h) > k:
            heapq.heappop(h)            # evict the least frequent survivor
    return [value for _, value in h]    # order within the answer is unspecified

O(n + m log k) time, O(m) space for the counter plus O(k) for the heap. Values are integers here so tuple comparison never falls through to a non-comparable payload — but if the payload were a list or a tree node you would push (c, i, payload) with a running i.

The follow-up they escalate to: "make it strictly better than O(n log n) even when k is close to m." Bucket sort answers it. Counts are bounded by n, so index an array of length n+1 by count and walk it downwards:

def topKFrequent_buckets(nums, k):
    counts = Counter(nums)
    buckets = [[] for _ in range(len(nums) + 1)]   # index = frequency
    for value, c in counts.items():
        buckets[c].append(value)
    out = []
    for c in range(len(buckets) - 1, 0, -1):       # walk high frequency first
        for value in buckets[c]:
            out.append(value)
            if len(out) == k:
                return out
    return out

O(n) time and O(n) space. This only works because counts are bounded by n — name that precondition out loud, because it is the thing that makes the trick legal.

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).

Memory hook — "opposite heap, weakest root"

The one thing that goes wrong under pressure is the direction flip, so peg it with a three-word phrase: opposite heap, weakest root.

• Opposite heap — the heap type is always the opposite of the word in the problem. "K largest" takes a min-heap; "K smallest" takes a max-heap. Whenever you write heapq and the problem said largest, your fingers should feel the mismatch and that mismatch is correct.
• Weakest root — the root of the capped heap is always the weakest member of your current winners, which is simultaneously the eviction target and, when the loop ends, the answer to "what is the K-th best".

The shape-of-the-code peg is four lines and never varies: push, check size, pop if over, done. If your loop has an if x > h[0] guard instead, that is a valid micro-optimisation but the same four beats.

For heapq itself the peg is "min-only, so negate twice" — once on the way in, once on the way out. Never negate only once; that is the bug that produces the right shape and the wrong sign.

What interviewers actually ask

Top-K shows up constantly because it is fast to state, has an obvious wrong answer (sort everything), and has a clean follow-up ladder into streaming.

• 215 · Kth Largest Element in an Array (Medium) — Amazon, Meta, Google. Probing whether you reach past sorting. Follow-up: "can you do better than O(n log k)?" — quickselect, O(n) average.
• 347 · Top K Frequent Elements (Medium) — Amazon, Meta, Google, Uber. Probing the count-then-heap composition. Follow-up: "your algorithm must be better than O(n log n)" — bucket sort by frequency gives O(n).
• 23 · Merge k Sorted Lists (Hard) — Amazon, Google, Meta, Microsoft. Probing the k-way merge and, critically, the tiebreaker index. Follow-up: compare against divide-and-conquer pairwise merging, which hits the same bound without a heap.
• 295 · Find Median from Data Stream (Hard) — Amazon, Google, Meta. Probing the two-heap balance invariant. Follow-up: "what if all numbers are in the range 0 to 100?" — a counting array beats both heaps.
• 973 · K Closest Points to Origin (Medium) — Amazon, Meta, Bloomberg. Probing top-K with a computed key. Follow-up: "why can you skip the square root?" — squared distance preserves ordering.
• 703 · Kth Largest Element in a Stream (Easy) — Amazon, Bloomberg. Probing the streaming form where the heap is class state, not a local.
• 621 · Task Scheduler (Medium) — Amazon, Meta. Probing a max-heap driving a greedy schedule rather than answering a top-K question directly.

The escalation to watch for: almost every one of these ends with "now do it on an infinite stream." That question is why the heap exists — sort-and-slice has no answer to it.

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.

Quiz

Practice queue

Grind these in order, easy to hard. The two marked as anchors are the ones the rest reduce to — schedule them for spaced repeat.

1. 703 · Kth Largest Element in a Stream — Easy. The size-K loop as persistent class state.
2. 1046 · Last Stone Weight — Easy. Max-heap by negation, nothing else.
3. 215 · Kth Largest Element in an Array — Medium. Anchor, +7d then +21d.
4. 973 · K Closest Points to Origin — Medium. Top-K on a computed key.
5. 347 · Top K Frequent Elements — Medium. Anchor, +7d then +21d. Solve it both ways.
6. 621 · Task Scheduler — Medium. Max-heap driving a greedy schedule.
7. 373 · K Pairs with Smallest Sums — Medium. Lazy expansion — never materialise all pairs.
8. 23 · Merge k Sorted Lists — Hard. The tiebreaker index lives or dies here.
9. 295 · Find Median from Data Stream — Hard. Two heaps; returns properly in L52.

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.