Python Priority Queue
Aug 6th, 2021 - written by Kimserey with .
Priority queues are useful to keep track of smallest elements in Python. A typical example would be to keep track of the smallest elements of a collection, for example first, second, third elements, we can simply keep popping out of the priority queue to get them. Python comes with a built in pirority queue via the library heapq. In today’s post, we will look at the main functionalities of heapq with examples.
Heap
For this post, we start first by importing all the functions we will be using:
1
from heapq import heapify, heappush, heappop, heappushpop
heapq uses a Heap to keep track of the items. A heap is a binary tree where the parent node is inferior or equal to its children. The heap can be flattened into an array where the children of a node k are located at 2*k + 1 for the left child and 2*k + 2 for the right child; hence h[k] <= h[2*k + 1] and h[k] <= h[2*k + 2].
Heapify
heapq functions work on an array inplace. This means that elements will be shifted within the array keeping the memory space constant. The first function heapify will heapify the array. Heapifying is the action of bubbling back up to the root the smallest of all elements.
1
2
3
4
5
>>> a = [10, 20, 2, 30, 3]
>>> heapify(a)
>>> a
>>> [2, 3, 10, 30, 20]
We can see that we started with an unsorted array, and after heapifying it, we get 2 moved to a[0]. If we start from an array with values already in it, it is important to execute a heapify so that the value at index 0 is ensured to be the smallest one.
Push
The next function is heappush which allows us to push new items onto the heap.
1
2
3
4
>>> heappush(a, 15)
>>> a
>>> [2, 3, 10, 30, 20, 15]
After pushing 15, we can see that the heap is still maintained. Internally, heapify is being called.
1
2
3
4
>>> heappush(a, 1)
>>> a
>>> [1, 3, 2, 30, 20, 15, 10]
After pushing 1, we can see that 1 was pushed to the top.
Pop
Popping an item from the heap will give us the lowest element:
1
2
3
4
5
6
7
8
9
10
11
12
13
>>> a = [1, 3, 2, 30, 20, 15, 10]
>>>: heappop(a)
>>> 1
>>> heappop(a)
>>> 2
>>> heappop(a)
>>> 3
>>> a
>>> [10, 15, 20, 30]
And we can see that as we pop elements, the heap bubbles up the lowest element, making it available to pop.
Push and Pop
A common scenario is to push a new element and pop from the heap in order to maintain the same length on the array.
To ease the usage, heapq also implements heappushpop:
1
2
3
4
5
6
7
8
>>> a
>>> [10, 15, 20, 30]
>>> heappushpop(a, 20)
>>> 10
>>> a
>>> [15, 20, 20, 30]
We can see that after push pop, 20 was added and 10 was popped out, the array length was maintained.
Usage with Tuples
We can also use heapq with tuple, this allows us to put a priority on types that wouldn’t otherwise have an explicit comparison.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
>>> a = [(10, "apple"), (3, "pear"), (30, "orange")]
>>> heapify(a)
>>> a
>>> [(3, 'pear'), (10, 'apple'), (30, 'orange')]
>>> heappop(a)
>>> (3, 'pear')
>>> heappop(a)
>>> (10, 'apple')
>>> heappop(a)
>>> (30, 'orange')
We can see that the items are coming out the queue in ascending order.
Max Heap
Lastly one aspect of heapq is that it is a min heap. Every parents are smaller than their children. There are cases where we want a max heap instead so that we can pop max items out of the heap. To achieve this with heapq while working with numbers, we can invert the sign, for example:
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
>>> a = [(10, "apple"), (3, "pear"), (30, "orange")]
>>> a = [(-p, v) for p, v in a]
>>> a
>>> [(-10, 'apple'), (-3, 'pear'), (-30, 'orange')]
>>> heapify(a)
>>> a
>>> [(-30, 'orange'), (-3, 'pear'), (-10, 'apple')]
>>> heappop(a)
>>> (-30, 'orange')
>>> heappop(a)
>>> (-10, 'apple')
>>> heappop(a)
>>> (-3, 'pear')
We can see that the items are then coming out of the queue in descending order. And that concludes today’s post! I hope you liked this post and I see you on the next one!