Priority Queues in Python. Heap Implementation of Priority Queues | by Sadrach Pierre, Ph.D. | Towards Data Science

A priority queue is a data structure that is similar to ordinary queues or stacks, but each element has an associated priority. Elements with high priority are served before lower priority elements. Priority queues are typically implemented using a heap data structure.
In this post, we will discuss the implementation of a priority queue in python using a heap data structure.
Let’s get started!
It is worth familiarizing ourselves with the python ‘heapq’ module before we build our ‘PriorityQueue’ class. To start let’s import the ‘heapq’ module:
import heapq
Suppose we have a list of historical prices for IBM:
ibm_prices = [108.68, 109.65, 121.01, 122.78, 120.16]
If we wanted to get the N lowest or highest prices we can use the ‘.nsmallest()’ and ‘.nlargest()’ methods respectively. To get the three lowest prices we can do the following:
print(heapq.nsmallest(3, ibm_prices))
And for the three highest prices:
print(heapq.nlargest(3, ibm_prices))
To make things a bit more interesting, suppose we have a portfolio of tech company stocks:
portfolio = [       {'name': 'IBM', 'shares': 200, 'price': 108.68},       {'name': 'AAPL', 'shares': 75, 'price': 277.97},       {'name': 'FB', 'shares': 40, 'price': 170.28},       {'name': 'HPQ', 'shares':125, 'price': 17.18},       {'name': 'UBER', 'shares': 50, 'price': 22.60},       {'name': 'TWTR', 'shares': 95, 'price': 29.29}]
We can use the ‘.nsmallest()’ and ‘.nlargest()’ to pull the cheapest and most expensive stocks respectively:
cheap_stocks = heapq.nsmallest(3, portfolio, key=lambda s: s['price'])expensive_stocks = heapq.nlargest(3, portfolio, key=lambda s: s['price'])
Let’s print the results:
print("Cheap Stocks: ", cheap_stocks)print("Expensive Stocks: ", expensive_stocks)
These function provide superior performance especially when the size of the N elements (largest or smallest) is small compared to the entire collection.
Now we are equipped to build our priority queue class. Let’s create a priority queue class using the ‘heapq’ module. First, we create the initialization method:
class PriorityQueue:    def __init__(self):        self._queue = []        self._index = 0
Let’s also define a method that will allow us to check if the queue is empty:
class PriorityQueue:    ...    def is_empty(self):        return not self._queue
Next, let’s define a method that will allow us to push objects into our priority queue. We employ the ‘heappush’ method which will take the priority and the item value:
class PriorityQueue:    ...    def push(self, item, priority):        heapq.heappush(self._queue, (priority, self._index, item))        self._index += 1
Finally, we will add a method that will let us remove elements from our priority queue, using the ‘heappop’ method:
class PriorityQueue:    ...    def pop(self):        return heapq.heappop(self._queue)[-1]
Next let’s define a class, called Stock , that we will use to demonstrate the use of our priority queue class. The class will have an ‘__init__’ method that will let us initialize values for stock price, shares and ticker:
class Stock:    def __init__(self, stock_ticker, stock_price, stock_share):        self.stock_ticker = stock_ticker        self.stock_price = stock_price        self.stock_share = stock_share
We can also define a method that allows us to print the representation of the class attributes:
class Stock:    ...    def __repr__(self):        return 'Stock({}, {}, {})'.format(self.stock_ticker ,   self.stock_price,  self.stock_share)
Now we are ready to initialize our priority queue and add items:
q = PriorityQueue()
Let’s add the stock attributes from our initial porfolio to our priority queue. We will prioritize cheap stocks:
q.push(Stock('IBM', 108.68, 200), 4)q.push(Stock('HPQ', 17.18, 125), 1)q.push(Stock('TWTR', 29.29, 95), 3)q.push(Stock('UBER', 22.6, 50), 2)q.push(Stock('AAPL', 277.97, 75), 6)q.push(Stock('FB', 170.28, 40), 5)
Now let’s print the items in our priority queue:
while not q.is_empty():    print(q.pop())
I’ll stop here but feel free to play around with the heap implementation of priority queues. Additionally, you can try building a priority queue with ordinary lists or linked lists. You can find some examples of other priority queue implementations in python here. I hope you found this post useful/interesting. The code from this post is available on GitHub. Thank you for reading!