heapq: Max-Heaps, Priorities and Top-K

Heaps & Top-K, lesson 2 of 3

heapq: Max-Heaps, Priorities and Top-K

Drive Python's min-heap like a pro and keep the k best items in O(n log k).

13 min, 0 of 3 activities solved

View cheatsheet

Getting Python ready… examples can run in a moment.

Python ships a heap in the heapq module. It works on an ordinary list, and it is a min-heap only.

Call Does Cost
heapq.heappush(h, x) add x O(log n)
heapq.heappop(h) remove + return min O(log n)
h[0] peek at min O(1)
heapq.heapify(h) list → heap in place O(n)
heapq.heappushpop(h, x) push, then pop O(log n)
heapq.heapreplace(h, x) pop, then push O(log n)
heapq.nlargest(k, it) k biggest, sorted O(n log k)
heapq.nsmallest(k, it) k smallest, sorted O(n log k)
Example

The basics

heapify changes the list in place and returns None, so never write h = heapq.heapify(h).

Need a max-heap? Negate. Push -x, and flip the sign again when you read it back. The smallest negative is the largest original.

Need priorities? Push tuples. Tuples compare element by element, so (priority, item) pops the lowest priority first. If two priorities tie, Python compares the next field, which crashes when that field can't be compared (dicts, objects, ListNodes). Put a unique counter in between: (priority, count, item). Ties then break by insertion order (FIFO) and the item is never compared.

Example

Max-heap and a tie-breaking counter

With the counter, "email" (pushed first) beats "backup" at the same priority, and the dicts are never compared.

Quiz

What does this print?

Top-k with a size-k heap

To find the k largest of n items, keep a min-heap of size k. Its root is the weakest of your current top k, the one to kick out:

  1. Push each item.
  2. If the heap grows past k, pop the root (the smallest).
  3. At the end the heap holds the k largest, and h[0] is the k-th largest.

Every heap operation costs O(log k), so the total is O(n log k) time and O(k) space. Compare:

  • sort everything: O(n log n) time, O(n) space
  • heapify all + pop k times: O(n + k log n), O(n)
  • size-k heap: O(n log k), O(k), and it works on a stream you can't hold in memory.

It feels backwards at first: largest k uses a min-heap. Flip both for the k smallest (a max-heap of size k, via negation).

Example

Top-3 on a stream, traced

sorted(h) is only for display. In real code heapq.heappushpop(h, x) does the push-then-pop in one step once the heap is full, and it's faster.

Quiz

You read a huge stream and must always know its 100 largest values. What do you keep?

Exercise

The k largest, best first

Write k_largest(nums, k) that returns the k largest values of nums in descending order. Use a min-heap that never holds more than k items (don't sort all of nums). Assume 0 <= k <= len(nums).

k_largest([5, 1, 9, 3, 7], 2)  ->  [9, 7]