heapq Min-Heap Recipes
Python's `heapq` module turns any list into a binary min-heap in place, supporting O(log n) push and pop. It is the priority-queue primitive that powers Dijkstra, Top-K, and merging sorted streams. This snippet covers the basic push and pop, the Top-K largest pattern using `nsmallest` and `nlargest`, and the merge of multiple sorted iterables in O(N log K).
587 views
17
import heapq
heap = []
for x in [5, 3, 7, 1, 9, 2]:
heapq.heappush(heap, x)
print(heap) # [1, 3, 2, 5, 9, 7] (heap-ordered, not sorted)
print(heap[0]) # 1 (smallest)
out = []
while heap:
out.append(heapq.heappop(heap))
print(out) # [1, 2, 3, 5, 7, 9]heapq operates directly on a Python list, treating it as a binary heap with heap[0] as the smallest element. heappush and heappop are both O(log n). Note that the heap is NOT sorted: it satisfies the heap invariant (parent <= children), but the in-memory order can look scrambled. Reading heap[0] gives an O(1) min query, which is the headline operation. To sort, repeatedly pop until empty (heap sort), which is O(n log n) total.
import heapq
scores = [42, 91, 17, 88, 23, 65, 99, 4, 76, 50]
print(heapq.nlargest(3, scores)) # [99, 91, 88]
print(heapq.nsmallest(3, scores)) # [4, 17, 23]
# With a key function (top-3 cheapest products)
products = [
{'name': 'pen', 'price': 2.50},
{'name': 'book', 'price': 14.00},
{'name': 'lamp', 'price': 45.00},
{'name': 'mug', 'price': 8.00},
]
print(heapq.nsmallest(2, products, key=lambda p: p['price']))heapq.nlargest(k, iter) and heapq.nsmallest(k, iter) return the K largest or smallest elements in O(n log k) time, which is asymptotically better than sorted(...)[:k] (O(n log n)) when k is much smaller than n. They both accept an optional key argument for projection, mirroring sorted. For very small K (1, 2) the constant factors favour min/max instead, but for K up to a few hundred the heap approach is the right default.
import heapq
shard1 = [1, 4, 9, 13]
shard2 = [2, 6, 10, 14]
shard3 = [3, 7, 11]
for x in heapq.merge(shard1, shard2, shard3):
print(x, end=' ')
print() # 1 2 3 4 6 7 9 10 11 13 14
# With a custom key
logs1 = [{'ts': 1, 'msg': 'a'}, {'ts': 4, 'msg': 'd'}]
logs2 = [{'ts': 2, 'msg': 'b'}, {'ts': 3, 'msg': 'c'}]
for entry in heapq.merge(logs1, logs2, key=lambda e: e['ts']):
print(entry)heapq.merge(*iters) lazily merges any number of pre-sorted iterables into one sorted stream. It uses a heap of size K (the number of inputs) so memory is bounded regardless of how long each input is, and total time is O(N log K) for N total items. This is the streaming-friendly answer to 'merge K sorted lists' and the building block for external sorts on data too big for RAM. The optional key argument lets you merge dicts or objects by an extracted field.