Introduction to Python Heapq Module
By: Vijini Mallawaarachchi
Tags: Programming, Computer Science, Algorithms, Data Structures, Data Science

Introduction to Python Heapq Module

A simple introduction on how to use Python’s heapq module

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.

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.

Image by Karen Arnold from Pixabay

Heapq Functions

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

— Push the value item into the heap heappop(heap) — Pop and return the smallest value from the 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

— Push the value item into the heap and return the smallest value from the heap heapify(x) — Convert the list x into a heap

— 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

Example Walkthrough of Heapq with Primitive Data Types

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.

How to Use Objects in Heapq?

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 strings

my_strings = ["write", "go", "run", "come"] # Initialising

sorted_strings = [] # Wrap strings and push to heap

for s in my_strings:

heapObj = DataWrap(s)

heapq.heappush(sorted_strings, heapObj) # Print the heap

for 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.

Implementing a Max Heap using Heapq

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 strings

my_strings = ["write", "go", "run", "come"] # Initialising

sorted_strings = [] # Wrap strings and push to heap

for s in my_strings:

heapObj = DataWrap(s)

heapq.heappush(sorted_strings, heapObj) # Print the heap

for 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.

Final Thoughts

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!

References

[1] heapq — Heap queue algorithm — Python 3.8.3rc1 documentation (https://docs.python.org/3/library/heapq.html)