The smallest thing, cheaply

A heap answers one question fast and everything else slowly, which is exactly the trade a lot of problems want.

Key idea

One question, answered fast

A heap keeps a collection in an arrangement where the smallest item is always immediately available. Adding an item and removing the smallest each cost O(log n), and looking at the smallest costs nothing.

It does not keep the collection sorted. Asking for the second smallest, or for any particular item, is not fast. That narrowness is the point: sorting everything costs O(n log n) up front, and a heap gives you the one thing you actually need for O(log n) per operation.

1import heapq2 3values = [5, 1, 8, 3]4heapq.heapify(values)          # rearrange in place, O(n)5print(values[0])               # the smallest, no work6 7heapq.heappush(values, 2)8print(heapq.heappop(values), heapq.heappop(values))9print(values)

Key idea

Python has no max-heap

heapq is a min-heap and there is no max-heap in the standard library. The standard workaround is to negate the values on the way in and negate them again on the way out.

For tuples, negate the field being ordered by. Doing this by hand each time is error prone, so keep the negation at the boundaries: negate when pushing, negate when popping, and never let a negated value escape into the rest of the code.

1import heapq2 3heap = []4for value in [5, 1, 8, 3]:5    heapq.heappush(heap, -value)      # negate going in6 7largest = -heapq.heappop(heap)        # negate coming out8second = -heapq.heappop(heap)9print(largest, second)

Tip

Ordering by something other than the value

Push tuples and the heap orders by the first element, breaking ties on the second, and so on. That is how you order by a computed key: put the key first and the payload after it.

The catch is that ties fall through to the next element, so every element in the tuple must be comparable. Two tuples whose first elements are equal will compare their second, and if that is a custom object with no ordering defined, Python raises. Inserting a unique counter as a tie-breaker before any non-comparable payload avoids it.

1import heapq2 3heap = []4for index, word in enumerate(["pear", "fig", "banana"]):5    heapq.heappush(heap, (len(word), index, word))    # length, then arrival6 7while heap:8    length, _, word = heapq.heappop(heap)9    print(length, word)

Predict the output· 5 min

Predict the output: reading a heap by index

Type the two lines this prints.

Python

1import heapq2 3values = [9, 7, 5, 3, 1]4heapq.heapify(values)5print(values[0], values[-1])6print(sorted(values) == values)

Printed output

Code lab· 25 min

Heap drills

k_smallest returns the k smallest values in ascending order. k_largest returns the k largest in descending order, using negation. merge_sorted_lists merges several already-sorted lists into one sorted list using a heap.

solution.py

Tests

print(k_smallest([5, 1, 8, 3], 2), k_smallest([2], 5))print(k_largest([5, 1, 8, 3], 2), k_largest([], 3))print(merge_sorted_lists([[1, 4, 5], [1, 3, 4], [2, 6]]))print(merge_sorted_lists([]))

Output

Run tests to see output.

1import heapq2 3heap = []4for index, word in enumerate(["pear", "fig", "banana"]):5 heapq.heappush(heap, (len(word), index, word)) # length, then arrival6 7while heap:8 length, _, word = heapq.heappop(heap)9 print(length, word)