Heaps and Priority Queues
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.
Gotcha
Printing the list will disappoint you
The final list above is [3, 5, 8], which happens to look sorted. On larger inputs it will not be, and code that reads a heap by index for anything other than position 0 is wrong.
The only guarantee is that position 0 holds the smallest. Everything else is an internal arrangement. If you need the items in order, pop them one at a time, which is O(n log n) and is just heapsort.
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.
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.
Tip
The convenience functions
heapq.nlargest(k, items) and heapq.nsmallest(k, items) do the bounded-heap job for you, and both accept a key function. They are the right answer in real code.
In an interview, name them, then write the loop if asked. Knowing that nlargest is O(n log k) rather than a sort is the part being tested.
Cost
The costs worth memorizing
Building a heap from an existing list with heapify is O(n), which is better than pushing n items one at a time at O(n log n). Push and pop are each O(log n). Peeking is O(1).
So the pattern that beats sorting is: heapify once if you have everything up front, then pop only as many as you need. Pulling k items out of a heap of n costs O(n + k log n), and when k is small that is much better than sorting.
Predict the output: reading a heap by index
Type the two lines this prints.
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.
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 the tests when you are ready.