Heaps

Data Structures in Practice

Chapter 6 ยท Heaps

A priority queue hands back the item that matters most, not the one that arrived first or last. A heap is the structure that does this cheaply: it keeps just enough order to find the smallest item instantly, and does not waste effort sorting the rest. Python provides it in the heapq module.

What was run, and what was not
Every snippet below was run and the output shown is what it printed. No timings are quoted, because they differ between machines; costs are described as patterns.

What a Heap Gives You

heapq works on an ordinary Python list. The rule it maintains is that the item at index 0 is always the smallest, and every item is no larger than its two "children" at positions 2i+1 and 2i+2. Nothing else is promised: the list as a whole is not sorted.

import heapq h = [] for n in [5, 1, 8, 3, 2, 9]: heapq.heappush(h, n) print(h) print(h[0]) out = [] while h: out.append(heapq.heappop(h)) print(out)
[1, 2, 8, 5, 3, 9] 1 [1, 2, 3, 5, 8, 9]

The printed list is a valid heap but is not in order (8 comes before 5). Only h[0] is guaranteed. Popping repeatedly produces everything in sorted order, because each pop re-establishes the rule for what is left. The picture of that list as a tree:

index: 0 1 2 3 4 5 value: [ 1, 2, 8, 5, 3, 9 ] 1 / \ 2 8 / \ / 5 3 9

Building and Reading a Heap

If you already have all the items, do not push them one by one. heapify rearranges a list in place in linear time.

data = [5, 1, 8, 3, 2, 9] heapq.heapify(data) print(data) print(heapq.heappop(data), data)
[1, 2, 8, 3, 5, 9] 1 [2, 3, 8, 9, 5]
OperationCallCost
Look at the smallesth[0]Constant
Add an itemheappush(h, x)Logarithmic
Remove the smallestheappop(h)Logarithmic
Turn a list into a heapheapify(h)Linear
Find an arbitrary valuex in hLinear (a heap is not searchable)
Fully sorted orderpop everythingn pops, so n log n
Never reorder a heap by hand
Do not call h.sort() on something you will use as a heap (a sorted list happens to be a valid heap, but you will soon be tempted to insert or append directly, which silently breaks the rule). Change a heap only through heappush, heappop, heapreplace and heapify.

Max-Heaps by Negation

heapq only offers a min-heap. For "largest first" with numbers, store the negatives and flip the sign on the way out.

nums = [5, 1, 8, 3] mh = [-n for n in nums] heapq.heapify(mh) print(-heapq.heappop(mh))
8

Priorities, Ties and Tuples

To attach a payload to a priority, push a tuple. Python compares tuples element by element, so the priority decides first and the payload only matters on a tie.

h = [] heapq.heappush(h, (2, "write")) heapq.heappush(h, (1, "deploy")) heapq.heappush(h, (2, "test")) print(heapq.heappop(h), heapq.heappop(h), heapq.heappop(h))
(1, 'deploy') (2, 'test') (2, 'write')

Two things went wrong quietly here. On a tie the payload strings were compared, so "test" beat "write" alphabetically rather than by arrival order. Worse, if the payloads cannot be compared at all, a tie raises an error:

h = [] try: heapq.heappush(h, (1, {"a": 1})) heapq.heappush(h, (1, {"b": 2})) except TypeError as e: print("TypeError", e)
TypeError '<' not supported between instances of 'dict' and 'dict'

The standard fix is a tie-breaker counter in the middle of the tuple: (priority, count, item). The count is unique and increasing, so comparison never reaches the item, and equal priorities come out in first-in, first-out order.

import itertools counter = itertools.count() h = [] for pr, name in [(2, "write"), (1, "deploy"), (2, "test"), (1, "backup")]: heapq.heappush(h, (pr, next(counter), name)) while h: pr, _, name = heapq.heappop(h) print(pr, name)
1 deploy 1 backup 2 write 2 test

Top k Without Sorting Everything

"The three highest scores" does not need a full sort. nlargest and nsmallest take a key function like sorted does.

scores = [41, 7, 93, 15, 88, 62, 3, 77] print(heapq.nlargest(3, scores), heapq.nsmallest(3, scores)) words = ["pear", "fig", "banana", "kiwi", "apple"] print(heapq.nlargest(2, words, key=len)) print(heapq.nsmallest(2, words, key=len))
[93, 88, 77] [3, 7, 15] ['banana', 'apple'] ['fig', 'pear']

The idea behind them is worth knowing, because it works on streams too. Keep a min-heap holding the k best items seen so far. Its root is the weakest of the group, so a new item only matters if it beats the root, and then heapreplace swaps them in one step. Memory stays at k items however long the input is.

def top_k(stream, k): h = [] for x in stream: if len(h) < k: heapq.heappush(h, x) elif x > h[0]: heapq.heapreplace(h, x) return sorted(h, reverse=True) print(top_k(scores, 3))
[93, 88, 77]
NeedBest tool
Smallest or largest single itemmin() / max()
A few best items (k small)heapq.nlargest / nsmallest
Most of the items, or all, in ordersorted()
Repeatedly "the best now" while items keep arrivinga heap

A Few More Tools

  • heapreplace(h, x) pops the smallest, then pushes x. The returned value may be larger than x.
  • heappushpop(h, x) pushes x first, then pops the smallest, so it returns x itself if x is already the smallest.
  • heapq.merge(*sorted_inputs) lazily merges several already-sorted inputs into one sorted stream.
print(heapq.heappushpop([4, 6], 1), heapq.heappushpop([4, 6], 9)) print(list(heapq.merge([1, 4, 9], [2, 3, 10], [0, 5])))
1 4 [0, 1, 2, 3, 4, 5, 9, 10]

queue.PriorityQueue

The queue module from Data Structures in Practice 4 also has PriorityQueue. It is a heap wrapped in locks so several threads can share it safely. It uses the same tuple convention, with the same tie-breaking trap.

import queue q = queue.PriorityQueue() q.put((2, "b")) q.put((1, "a")) print(q.get(), q.qsize())
(1, 'a') 1

Use plain heapq in single-threaded code (it is lighter and gives you h[0] to peek) and PriorityQueue when threads share the queue.

Scheduling, and Cancelling Lazily

Schedulers, "next event" simulations, Dijkstra's shortest paths (Data Structures in Practice 8) and merging many sorted files are all the same pattern: repeatedly take the most urgent item while new ones keep arriving. One catch is that you cannot cheaply remove an item from the middle of a heap. The usual answer is lazy deletion: mark the entry as cancelled, and discard it when it eventually reaches the top. Exercise 3 builds this.

Heap, sorted list or set?
If you only ever need the best item and items keep arriving, a heap is right. If you need to search, remove arbitrary items or walk everything in order, a heap is the wrong shape: use a sorted list (with bisect), a dict or a tree (Data Structures in Practice 7).

Challenges

Exercise 1

Write a class KthLargest(k) with an add(x) method that accepts a stream of numbers and returns the current k-th largest value (or None until k values have arrived). Use a min-heap of at most k items. Feed it 4, 5, 8, 2, 10, 1, 9 with k = 3 and print each result.

๐Ÿ“„ View solution
Exercise 2

Write heap_sort(items) using only heapify and heappop (it must not change the input), and merge_sorted(*lists) that merges any number of sorted lists by keeping one candidate per list in a heap. Check the second against heapq.merge, and test it with an empty list and with no arguments.

๐Ÿ“„ View solution
Exercise 3

Build a Scheduler class with add(priority, name) (returning a ticket), cancel(ticket) and next_task() (returning None when empty). Lower numbers run first, equal priorities run in the order added, and cancelled tasks never run. Cancel lazily. Show it working with five tasks, one of them cancelled.

๐Ÿ“„ View solution

Quick Reference

  • Heap: a list where h[0] is always the smallest; the rest is only partly ordered
  • Costs: peek constant; push and pop logarithmic; heapify linear
  • Max-heap: store negatives (numbers) and flip the sign back
  • Payloads: push (priority, count, item) so ties are first-in, first-out and items are never compared
  • Top k: nlargest / nsmallest, or a size-k min-heap for streams
  • Others: heapreplace, heappushpop, heapq.merge, queue.PriorityQueue for threads
  • Cancelling: mark and skip when it surfaces (lazy deletion)
  • Not for: searching, arbitrary removal or in-order walking
Coming next
Data Structures in Practice 7 covers trees: binary trees, binary search trees, traversal and balance. A heap is a tree squeezed into a list; next the tree gets its own nodes.