Heap
Heaps have special properties that make it perfect for solving problems involving finding the min or max of something:
- It will always give you either its largest or smallest element
- It takes O(log(n)) time to add an item to a heap
- It takes O(1) time to remove the top element in a heap
Examples of questions that we could use heaps to solve are: find the shortest distance between two nodes in a weighted graph, keep track of the n largest digits.
In Python, a heap is a min-heap by default: the top of the heap is always the smallest item. (If you want the largest instead, see the max-heap section below.)
Creating a heap with heapify()
In Python, the two popular ways to create a heap is via either a PriorityQueueor heapq. The difference between them is heapq is a library that performs functions on a heap (which is a list), while PriorityQueue is a pre-built class.
We'll use heapq for these two reasons:
heapqallows you to start the heap with existing itemsheapqallows you to print out the entire heap
Note
heapq also tends to be a little faster because it is not thread-safe. However, interviewers do not care if you pick a fast library or not, so just use your favorite.
Our heap can simply be represented with an empty list:
Or we can initialized the list with existing items, and then use heapify() to rearrange the list into the heap:
Note
heapify() does not sort the heap, it only ensures the first/top-most item is the smallest item. heapify() also takes O(n) time to arrange all the items in the heap.
If the items in the heap are tuples, heapify() use the first value in the tuple to arrange by:
Adding an item to the heap
O(log(n))Add an item to the heap via the heappush() function:
Removing an item from the heap
O(log(n))To remove an item from a heap, use heappop() . This returns the item with the lowest value in the heap:
Peek at the top of the heap
O(1)There's no peek() function in heapq, and you don't need one. Since the heap is a list, you can look at index 0 of the list to peek at the smallest value in the heap:
Making a max-heap in Python
If you want to create a max-heap (a heap that will return the maximum value), use a negative priority when you push and flip the sign back when you pop.
How heapq compares tuples
If tuples are stored in the heap, heapq will attempt to arrange the items based on the first values of the tuples, and then the second values if the first values are equivalent. This is especially helpful to keep in mind when you need to store extra data for solutions, like finding the shortest path in a weighted graph.
Note
Watch out for a TypeError. If two first values tie and the second values can't be compared (two dicts, two custom objects), Python raises TypeError: '<' not supported between instances. The standard fix is to put a unique counter in the middle: (priority, count, item). The counter breaks every tie, so the third value is never compared.
Can you heapify a dictionary?
Not directly. heapify() only works on lists, so convert the dictionary first. list(d.items()) gives you (key, value) tuples, which heapify by key:
To order by value instead, build the tuples yourself: [(v, k) for k, v in d.items()].