Learning Priority Queue in Python

I was recently solving a leet code question that looks efficient with a priority queue. Now, coming from a C# background, where PriorityQueue is an out-of-the-box feature, makes me wonder if Python contains something similar.

This landed me on a Python module called Heapq, which provides a priority queue-like feature for Python.

Before scratching the library, let's look at some basics.

What is a Priority Queue?

A priority queue is a special type of queue in the data structure where each element is associated with a priority. In this queue, elements with higher priority are dequeued before the elements with lower priority. If two elements carry the same priority, they are served as per their order in the queue.

Now, the one above is a definition from Google, but to try to understand it better, let's compare it with a normal queue.

A Queue follows the principle of FIFO(First In, First Out), which means that the element that was added first to the collection will be fetched out first. This is irrespective of which value is higher or lower.

A Priority Queue, on the other hand, adds the elements as they are provided but assigns a priority to each element, usually on the basis of the value being smaller or bigger than other values. If two same values are added to the priority queue then a higher priority is assigned to the element which was added earlier. When we try to fetch/pop value out of the priority queue, the value with higher priority is removed first.

Let's try to implement our own Priority queue in Python.

class PriorityQueue:
    def __init__(self):
        self.nums = [6, 7, 9, 4, 3, 5, 8, 10, 1]

    def add(self, val):
        self.nums.append(val)

    def isEmpty(self):
        return len(self.nums) == 0

    def delete(self):
        try:
            maxVal = 0
            for i in range(len(self.nums)):
                if self.nums[i] < self.nums[maxVal]:
                    maxVal = i

            item = self.nums[maxVal]
            self.nums.remove(item)
            return item
        except IndexError:
            print()
            exit()
myQueue = PriorityQueue()
myQueue.add(12)
myQueue.add(1)
myQueue.add(14)
myQueue.add(7)
print(myQueue.nums)   # output = [6, 7, 9, 4, 3, 5, 8, 10, 1, 12, 1, 14, 7]
print("=======")
while not myQueue.isEmpty():
    print(myQueue.delete())
# output
# 1
# 1
# 3
# 4
# 5
# 6
# 7
# 7
# 8
# 9
# 10
# 12
# 14

Although this works fine, as you can see, the delete() method has a time complexity of O(n), which means it always goes through each element of the array to get and remove the desired value. This approach is not efficient with large data sets.

Heapq

The Heap data structure is mainly used to represent a priority queue. In Python, it is available using the “heapq” module. The property of this data structure in Python is that each time, the smallest heap element is popped(min-heap). Whenever elements are pushed or popped, heap structure is maintained. The heap[0] element also returns the smallest element each time.

This means that this module provides a property that the elements will be added as they are provided, but when trying to pop a value, it returns the smallest values in the most efficient manner.

import heapq
class HeapQExample:
    def __init__(self):
        self.nums = [6, 7, 9, 4, 3, 5, 8, 10, 1]
        heapq.heapify(self.nums)

    def displayHeap(self):
        for item in self.nums:
            print(item)

    def isEmpty(self):
        return len(self.nums) == 0

    def addToHeap(self, value):
        heapq.heappush(self.nums, value)

    def getSmallesValue(self):
        return heapq.heappop(self.nums)

    def popAllElements(self):
        while not self.isEmpty():
            print(self.getSmallesValue())

    def getNSmallestItems(self, n):
        return heapq.nsmallest(n, self.nums)

    def getNLargestItems(self, n):
        return heapq.nlargest(n, self.nums)

    def pushPopExample(self, val):
        return heapq.heappushpop(self.nums, val)

    def replaceExample(self, val):
        return heapq.heapreplace(self.nums, val)

hqe = HeapQExample()

hqe.displayHeap()
hqe.addToHeap(11)
print('======')
hqe.displayHeap()
print('=======')
item = hqe.getSmallesValue()
print(item)    # output = 1

print(hqe.getNSmallestItems(3)) # output = [1, 3, 4]
print(hqe.getNLargestItems(3)) # output = [11, 10, 9]
print(hqe.pushPopExample(-1))   # output = -1
print(hqe.replaceExample(-1))  # output = 1

As you can see in the result, the values are added in the same order as they are provided(confirmed from the display method), but then dequeued/popped, and the smaller value is fetched.

A brief definition of each of the methods that are referred to in the code snippet above.

  • heappush(heap, ele): Used to insert the element mentioned in its arguments into a heap. The order is adjusted so that heap structure is maintained.
  • heappop(heap): Used to remove and return the smallest element from the heap. The order is adjusted, so that heap structure is maintained.
  • heappushpop(heap, ele): This function combines the functioning of both push and pop operations in one statement, increasing efficiency. Heap order is maintained after this operation.
  • heapreplace(heap, ele): This function also inserts and pops elements in one statement, but it is different from the above function. In this, the element is first popped, then the element is pushed. i.e., the value larger than the pushed value can be returned. heapreplace() returns the smallest value originally in the heap regardless of the pushed element as opposed to heappushpop().
  • nlargest(k, iterable, key = fun): This function is used to return the k largest elements from the iterable specified and satisfy the key if mentioned.
  • nsmallest(k, iterable, key = fun): This function is used to return the k smallest elements from the iterable specified and satisfy the key if mentioned.

I hope this article helps you.

Up Next
    Ebook Download
    View all
    Learn
    View all