Introduction to Python Heapq Module | by Vijini Mallawaarachchi | Towards Data Science

A Heap is a special case of a binary tree where the parent nodes are compared to their children with their values and are arranged accordingly. If you have come across my previous article titled 8 Common Data Structures every Programmer must know that there are two types of heaps; min heap and max heap.
towardsdatascience.com
The heapq module of python implements the heap queue algorithm. It uses the min heap where the key of the parent is less than or equal to those of its children. In this article, I will introduce the python heapq module and walk you through some examples of how to use heapq with primitive data types and objects with complex data.
Assuming that you know how the heap data structure works, let’s see what functions are provided by Python’s heapq model.
heappush(heap, item) — Push the value item into the heap
heappop(heap) — Pop and return the smallest value from the heap
heappushpop(heap, item) — Push the value item into the heap and return the smallest value from the heap
heapify(x) — Convert the list x into a heap
heapreplace(heap, item) — Pop and return the smallest value from the heap then push the value item into the heap
Let’s see some examples where we use different heapq functions. First of all, you have to import the heapq module.
import heapq
Consider the example list a as given below.
>>> a = [52, 94, 13, 77, 41]
If we heapify this list, the result will be as follows. Note that heapifying is done in-place.
>>> heapq.heapify(a)>>> print(a)[13, 41, 52, 77, 94]
Note that the 0th index contains the smallest value out of all the values, which is 13.
Let’s push the value 10 to our heap.
>>> heapq.heappush(a,10)>>> print(a)[10, 41, 13, 77, 94, 52]
Note that 10 is added and since 10 is the smallest value out of the available values, it is in the 0th index.
Now let’s pop from our heap.
>>> print(heapq.heappop(a))10>>> print(a)[13, 41, 52, 77, 94]
When we pop from our heap, it will remove the smallest value from the heap and return it. Now the value 10 is no longer in our heap.
Let’s see how the heappushpop() functions works. Let’s heappushpop the value 28.
>>> print(heapq.heappushpop(a,28))13>>> print(a)[28, 41, 52, 77, 94]
You can see that 28 is pushed to the heap and the smallest value 13 is popped from the heap.
Now let’s try the heapreplace() function. Let’s heapreplace the value 3 in to the heap.
>>> print(heapq.heapreplace(a,3))28>>> print(a)[3, 41, 52, 77, 94]
You can see that the initially smallest value 28 is popped and then our new value 3 is pushed. The 0th index in the new heap value 3. Note the difference in the order of push and pop actions in heappushpop() and heapreplace() functions.
In the previous example, I have explained how to use heapq functions with primitive data types such as integers. Similarly, we can use objects with heapq functions to order complex data such as tuples or even strings. For this, we need to have a wrapper class according to our scenario. Consider a case where you want to store strings and order them by the length of the strings, shortest to longest. Our wrapper class will look as follows.
class DataWrap:        def __init__(self, data):        self.data = data          def __lt__(self, other):        return len(self.data) < len(other.data)
The __lt__ function (it is the operator overloading function for comparison operators; >, ≥, ==, < and ≤) will contain the logic to compare the lengths of strings. Now let’s try to make a heap with some strings.
# Create list of stringsmy_strings = ["write", "go", "run", "come"]# Initialisingsorted_strings = []# Wrap strings and push to heapfor s in my_strings:    heapObj = DataWrap(s)    heapq.heappush(sorted_strings, heapObj)# Print the heapfor myObj in sorted_strings:    print(myObj.data, end="\t")
The output from printing items from the heap will be as follows.
go	come	run	write
Note that the shortest word go is in the front of the heap. Now you can try out other heapq functions by wrapping strings.
We can easily implement a max heap using heapq. All you have to do it change the comparison operator in the __lt__ function to order the largest value to the front. Let’s try the previous example with strings and their lengths.
class DataWrap:        def __init__(self, data):        self.data = data          def __lt__(self, other):        return len(self.data) > len(other.data)# Create list of stringsmy_strings = ["write", "go", "run", "come"]# Initialisingsorted_strings = []# Wrap strings and push to heapfor s in my_strings:    heapObj = DataWrap(s)    heapq.heappush(sorted_strings, heapObj)# Print the heapfor myObj in sorted_strings:    print(myObj.data, end="\t")
Note how the length comparison has changed in the __lt__ function of the DataWrap class. The output from printing items from this heap will be as follows.
write	come	run	go
Note that now the longest word write is in the front of the heap.
I hope you found this article informative and useful during implementations with the heapq module. Please feel free to play around with the code provided.
Thank you for reading!
Cheers!
[1] heapq — Heap queue algorithm — Python 3.8.3rc1 documentation (https://docs.python.org/3/library/heapq.html)