⛰️ Heaps & Top-K
Need the min/max over and over while data changes? A heap gives it in O(1) and updates in O(log n).
Recognize the pattern
"k largest / k smallest / k-th largest / k closest" → size-k heap (min-heap for largest, max-heap for smallest): O(n log k)
Repeatedly take the best item, change it, put it back (stones, tasks) → max-heap of values
Merge k sorted lists / arrays / rows → heap of the k current heads: O(N log k)
Median (or any middle split) of a growing stream → two heaps: max-heap low half, min-heap high half
Simulate events in time order (rooms, CPUs, cooldowns) → min-heap keyed by time
Data arrives as a stream you can't store or re-sort each time → heap instead of sorting
heapq basics (min-heap)
import heapq
h: list[int] = []
heapq.heappush(h, 5) # O(log n)
smallest = h[0] # peek, O(1)
smallest = heapq.heappop(h) # O(log n)
nums = [4, 1, 7]
heapq.heapify(nums) # O(n), in place, returns None
top = heapq.heappushpop(nums, 3) # push then pop
top = heapq.heapreplace(nums, 9) # pop then push
heapq.nlargest(2, nums) # [9, 7]
heapq.nsmallest(2, nums, key=abs)
Max-heap and priorities
Negate for a max-heap. Use tuples for priorities and add a counter so ties never compare the payload.
import heapq
from itertools import count
maxh: list[int] = []
heapq.heappush(maxh, -10)
largest = -maxh[0]
pq: list[tuple] = []
tie = count()
heapq.heappush(pq, (2, next(tie), {"job": "a"}))
pri, _, item = heapq.heappop(pq)
Top-k with a size-k heap
import heapq
def k_largest(nums: list[int], k: int) -> list[int]:
h: list[int] = [] # min-heap of size k
for x in nums:
if len(h) < k:
heapq.heappush(h, x)
elif x > h[0]: # beats the weakest
heapq.heapreplace(h, x)
return h # h[0] = k-th largest
# k smallest: same idea with a max-heap (push -x)
K-way merge
import heapq
def merge_k(lists: list[list[int]]) -> list[int]:
h = [(l[0], i, 0) for i, l in enumerate(lists) if l]
heapq.heapify(h)
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
Two heaps (running median)
import heapq
low: list[int] = [] # max-heap (negated), size n or n+1
high: list[int] = [] # min-heap
def add(x: int) -> None:
heapq.heappush(low, -x)
heapq.heappush(high, -heapq.heappop(low))
if len(high) > len(low):
heapq.heappush(low, -heapq.heappop(high))
def median() -> float:
if len(low) > len(high):
return -low[0]
return (-low[0] + high[0]) / 2
Scheduling and lazy deletion
- Event order: min-heap of
(time, ...); pop the next event, push what it triggers. - Ready vs cooling: max-heap of ready work plus a
queue of
(ready_time, work); move items back when their time comes. - Can't delete from the middle: mark entries stale
and skip them when popped
(
while h and h[0] in stale: heappop(h)).
Operation costs
| Operation | Time | Space |
|---|---|---|
| push / pop | O(log n) | — |
| peek h[0] | O(1) | — |
| heapify a list | O(n) | O(1) |
| heappushpop / heapreplace | O(log n) | — |
| top-k with size-k heap / nlargest | O(n log k) | O(k) |
| k-way merge of N items | O(N log k) | O(k) |
| running median: add / find | O(log n) / O(1) | O(n) |
| remove arbitrary item | O(n) (or lazy deletion) | — |
Watch for
heapq is a min-heap only: negate values for a max-heap and negate again when you read them.
The heap list is not sorted: only h[0] is guaranteed. Don't index h[1] or h[-1] expecting the 2nd smallest or the max.
heapify works in place and returns None:
h = heapq.heapify(h)sets h to None.Tuples with equal priorities compare the next field: add a counter (priority, count, item) when items aren't comparable.
Top-k largest uses a MIN-heap of size k (and vice versa); the wrong heap evicts the best items.
Editing a value inside the heap list breaks the heap rule: push a new entry and lazily skip the stale one.