Implementing a Priority Queue in Python

Technical writer

Python

10.07.2024

Reading time: 6 min

A priority queue is a special type of queue where each element is associated with a priority. In a priority queue, elements are dequeued in order of their priority, rather than their insertion order. This makes priority queues a powerful data structure for scenarios where elements with higher priorities need to be processed first.

Importance of Priority Queues in Programming

Priority queues are widely used in various programming applications, including:

Task scheduling in operating systems
Managing the order of events in simulations
Implementing Dijkstra's algorithm for finding the shortest path in graphs
Managing a to-do list where tasks have different levels of urgency

Implementing a Priority Queue Using heapq

Python provides a built-in library called heapq that can be used to implement priority queues. The heapq module offers an efficient way to maintain a heap, which is a binary tree where the parent node is always smaller than or equal to its child nodes (min-heap).

Basic Operations: Push, Pop, Peek

Let's explore the basic operations of a priority queue using the heapq module.

Pushing an Element

To add an element to the priority queue, use the heappush function:

import heapq

# Create an empty priority queue
priority_queue = []

# Add elements to the priority queue
heapq.heappush(priority_queue, (2, 'task 2'))
heapq.heappush(priority_queue, (1, 'task 1'))
heapq.heappush(priority_queue, (3, 'task 3'))

print(priority_queue)

Output:

Popping an Element

To remove and return the smallest element from the priority queue, use the heappop function:

smallest_task = heapq.heappop(priority_queue)
print(smallest_task)
print(priority_queue)

Output:

Peeking at the Smallest Element

To look at the smallest element without removing it, simply access the first element of the list:

# Peek at the smallest element
smallest_task = priority_queue[0]
print(smallest_task)

Output:

Example Use Cases for Priority Queues

Let's look at different use cases for priority queues.

Task Scheduling

In task scheduling, tasks with higher priorities should be executed before those with lower priorities. A priority queue can manage the order of task execution efficiently.

import heapq

tasks = []

# Add tasks with priorities
heapq.heappush(tasks, (1, 'write report'))
heapq.heappush(tasks, (3, 'email team'))
heapq.heappush(tasks, (2, 'prepare presentation'))

while tasks:
    priority, task = heapq.heappop(tasks)
    print(f'Executing: {task}')

Output:

Dijkstra's Algorithm

Dijkstra's algorithm finds the shortest path from a source node to all other nodes in a graph. A priority queue is used to select the next node to process based on the shortest known distance.

import heapq

def dijkstra(graph, start):
    priority_queue = []
    heapq.heappush(priority_queue, (0, start))
    distances = {node: float('inf') for node in graph}
    distances[start] = 0

    while priority_queue:
        current_distance, current_node = heapq.heappop(priority_queue)

        if current_distance > distances[current_node]:
            continue

        for neighbor, weight in graph[current_node]:
            distance = current_distance + weight
            if distance < distances[neighbor]:
                distances[neighbor] = distance
                heapq.heappush(priority_queue, (distance, neighbor))

    return distances

graph = {
    'A': [('B', 1), ('C', 4)],
    'B': [('A', 1), ('C', 2), ('D', 5)],
    'C': [('A', 4), ('B', 2), ('D', 1)],
    'D': [('B', 5), ('C', 1)],
}

distances = dijkstra(graph, 'A')
print(distances)

Output:

Advanced Priority Queue Operations

In this section, we will look closer at advanced usage of priority queues.

Updating an Element's Priority

Updating an element's priority involves removing the element and adding it again with the new priority. This can be inefficient, but it's a necessary step since heapq does not support direct priority updates.

import heapq

priority_queue = [(2, 'task 2'), (1, 'task 1'), (3, 'task 3')]
heapq.heapify(priority_queue)

# Remove task 2 and add it with a new priority
priority_queue = [(p, t) for p, t in priority_queue if t != 'task 2']
heapq.heapify(priority_queue)
heapq.heappush(priority_queue, (4, 'task 2'))

print(priority_queue)

Output:

Merging Two Priority Queues

To merge two priority queues, use the heapq.merge function, which returns an iterator over the sorted values.

import heapq

pq1 = [(1, 'task 1'), (3, 'task 3')]
pq2 = [(2, 'task 2'), (4, 'task 4')]

merged_pq = list(heapq.merge(pq1, pq2))
print(merged_pq)

Output:

Practical Code Examples

Priority Queue with Custom Objects

You can use custom objects in a priority queue by implementing comparison methods.

import heapq

class Task:
    def __init__(self, priority, description):
        self.priority = priority
        self.description = description

    def __lt__(self, other):
        return self.priority < other.priority

priority_queue = []
heapq.heappush(priority_queue, Task(2, 'task 2'))
heapq.heappush(priority_queue, Task(1, 'task 1'))
heapq.heappush(priority_queue, Task(3, 'task 3'))

while priority_queue:
    task = heapq.heappop(priority_queue)
    print(task.description)

Output:

Using Priority Queue in a Multi-threaded Environment

Python's queue.PriorityQueue is a thread-safe priority queue that can be used in multi-threaded applications.

import threading
import queue

def worker(task_queue):
    while not task_queue.empty():
        priority, task = task_queue.get()
        print(f'Processing {task}')
        task_queue.task_done()

task_queue = queue.PriorityQueue()
task_queue.put((2, 'task 2'))
task_queue.put((1, 'task 1'))
task_queue.put((3, 'task 3'))

threads = []
for i in range(3):
    thread = threading.Thread(target=worker, args=(task_queue,))
    thread.start()
    threads.append(thread)

for thread in threads:
    thread.join()

task_queue.join()

Output:

Conclusion

Priority queues are essential for efficiently managing tasks and resources based on priority. Python's heapq and queue.PriorityQueue modules provide powerful tools to implement and manipulate priority queues. Whether for simple tasks or complex algorithms like Dijkstra's, understanding how to use priority queues effectively can significantly enhance your programming toolkit.

If you want to build a web service using Python, you can rent a cloud server at competitive prices with Hostman.

For further reading and additional resources, consider the following:

Python

10.07.2024

Reading time: 6 min

Similar

Python

Implementing a Priority Queue in Python

Importance of Priority Queues in Programming

Implementing a Priority Queue Using heapq

Basic Operations: Push, Pop, Peek

Pushing an Element

Popping an Element

Peeking at the Smallest Element

Example Use Cases for Priority Queues

Task Scheduling

Dijkstra's Algorithm

Advanced Priority Queue Operations

Updating an Element's Priority

Merging Two Priority Queues

Practical Code Examples

Priority Queue with Custom Objects

Using Priority Queue in a Multi-threaded Environment

Conclusion

Similar

Python Static Method

Input in Python

Operators in Python

Do you have questions, comments, or concerns?

Do you have questions,
comments, or concerns?