Heaps / Top-K / Two Heaps — Coding Interview Patterns

When you need the biggest or smallest few — but not a full sort — a heap is the tool. It keeps the extreme element instantly available, which solves a whole family of “top-K” and “streaming” problems efficiently.

The idea

A heap keeps the minimum (min-heap) or maximum (max-heap) at the top in O(1), and supports insert/remove in O(log n) — all without sorting the rest. Python’s heapq is a min-heap; for a max-heap, push negatives (or use a wrapper). That partial order is exactly enough for “give me the extreme, repeatedly.”

The signature Top-K trick

Counter-intuitively, to find the k largest elements you keep a min-heap of size k. Push each number; whenever the heap grows past k, pop the smallest. Whatever survives is the top k, and the heap’s root is the k-th largest:

import heapq
def kth_largest(nums, k):
    heap = []
    for n in nums:
        heapq.heappush(heap, n)
        if len(heap) > k:
            heapq.heappop(heap)      # drop the smallest; keep the k biggest
    return heap[0]                   # root = k-th largest
# kth_largest([3,2,1,5,6,4], 2) -> 5

This is O(n log k) time and O(k) space — better than sorting’s O(n log n) when k is small, and it works on a stream (you never need all n at once). The same shape gives top-k frequent (heap by count) and k closest points (heap by distance).

Worked example: merge k sorted lists

Keep a heap of the current head of each list; repeatedly pop the smallest and push its successor — O(N log k) for N total elements across k lists:

import heapq
def merge_k_lists(lists):
    heap = [(l[0], i, 0) for i, l in enumerate(lists) if l]   # (value, list idx, pos)
    heapq.heapify(heap)
    out = []
    while heap:
        val, li, pos = heapq.heappop(heap)
        out.append(val)
        if pos + 1 < len(lists[li]):
            heapq.heappush(heap, (lists[li][pos+1], li, pos+1))
    return out
# merge_k_lists([[1,4,5],[1,3,4],[2,6]]) -> [1,1,2,3,4,4,5,6]

Two heaps: median of a data stream

To get the running median as numbers arrive, keep two heaps — a max-heap of the lower half and a min-heap of the upper half — balanced so their sizes differ by at most one. The median is then the top(s):

import heapq
class MedianFinder:
    def __init__(self):
        self.low = []    # max-heap (store negatives)
        self.high = []   # min-heap
    def add(self, x):
        heapq.heappush(self.low, -x)
        heapq.heappush(self.high, -heapq.heappop(self.low))   # balance value
        if len(self.high) > len(self.low):                    # balance size
            heapq.heappush(self.low, -heapq.heappop(self.high))
    def median(self):
        if len(self.low) > len(self.high): return -self.low[0]
        return (-self.low[0] + self.high[0]) / 2

Step-by-step walkthrough

Trace Kth Largest Element on [3, 2, 1, 5, 6, 4] with k = 2. Keep a min-heap of size k: whenever it grows past k, evict the smallest. What survives are the k largest, and the heap’s root is the k-th largest.

x	operation	heap (sorted view)	root
3	push 3	[3]	3
2	push 2	[2, 3]	2
1	push 1, size > 2 → pop 1	[2, 3]	2
5	push 5, size > 2 → pop 2	[3, 5]	3
6	push 6, size > 2 → pop 3	[5, 6]	5
4	push 4, size > 2 → pop 4	[5, 6]	5

The root ends at 5, the 2nd largest. Each step is O(log k), so the whole pass is O(n log k) — better than sorting when k is small.

When to reach for it

Signals: “top K” / “k-th largest or smallest” / “k closest”, “merge k sorted …”, “median of a stream”, “schedule tasks by frequency” (task scheduler), or “process the next-earliest/most-urgent thing”. Any time you repeatedly need the current extreme without sorting everything.

Pitfalls

Wrong heap direction for Top-K. k largest → min-heap of size k (so you can cheaply drop the smallest). This trips everyone; remember the inversion.
Python has no max-heap. Push -value (and negate on read), or heapify tuples. Don’t forget to flip the sign back.
Using a heap when sorting is simpler. If you need all elements sorted (not just top k) and have them all upfront, a plain sort is simpler and the same big-O. The heap wins for streaming or small k.

Recap

A heap gives the min/max in O(1) and insert/remove in O(log n) — ideal for top-K and streaming problems. The signature trick: min-heap of size k for the k largest (O(n log k)). Merge k sorted keeps a heap of current heads; running median balances two heaps. In Python, heapq is a min-heap — push negatives for a max-heap. Signal: top/k-th, merge k, median of a stream, by-frequency.

Practice: 11 solved problems on heaps & top-K → — cover each one, solve it yourself, then check the tested code.

Next: Lesson 15 — Backtracking.