60. Using heapq for Priority Queues

The heapq module in Python provides a way to implement efficient priority queues using heaps. A heap is a specialized tree-based data structure that satisfies the heap property, which ensures that the smallest element is always at the root. Here are several examples of using heapq for priority queues:

1. Basic Priority Queue Implementation

This is the basic implementation of a priority queue using heapq where the smallest element has the highest priority.

Copy

import heapq

# Create an empty list to represent the priority queue
pq = []

# Add items to the queue using heapq.heappush()
heapq.heappush(pq, (2, 'task 2'))
heapq.heappush(pq, (1, 'task 1'))
heapq.heappush(pq, (3, 'task 3'))

# Pop items from the queue using heapq.heappop()
while pq:
    priority, task = heapq.heappop(pq)
    print(f"Task: {task}, Priority: {priority}")

Output:

Copy

Task: task 1, Priority: 1
Task: task 2, Priority: 2
Task: task 3, Priority: 3

In this example, the heappush function is used to insert tasks with their associated priority, and heappop is used to remove and return the task with the highest priority (smallest value).

---

2. Priority Queue with Custom Priorities

In this example, we implement a priority queue where we control the order by using custom priorities (e.g., a negative priority value to invert the order).

Copy

import heapq

# Create an empty list to represent the priority queue
pq = []

# Add items to the queue with custom priorities (negative for highest priority)
heapq.heappush(pq, (-2, 'task 2'))  # Priority 2
heapq.heappush(pq, (-1, 'task 1'))  # Priority 1
heapq.heappush(pq, (-3, 'task 3'))  # Priority 3

# Pop items in order of highest priority
while pq:
    priority, task = heapq.heappop(pq)
    print(f"Task: {task}, Priority: {-priority}")

Output:

Copy

Task: task 3, Priority: 3
Task: task 2, Priority: 2
Task: task 1, Priority: 1

By negating the priority, we ensure that the heapq module uses the smallest value as the highest priority, effectively inverting the order.

---

3. Using heapq.nlargest() for Top-N Items

You can use heapq.nlargest() to efficiently find the n largest items in a collection, which is a common use case in priority queues.

Copy

import heapq

# Sample data
data = [(1, 'task 1'), (3, 'task 3'), (2, 'task 2')]

# Get the two largest items by priority
top_tasks = heapq.nlargest(2, data, key=lambda x: x[0])

for task in top_tasks:
    print(f"Task: {task[1]}, Priority: {task[0]}")

Output:

Copy

Task: task 3, Priority: 3
Task: task 2, Priority: 2

In this example, nlargest retrieves the top 2 tasks with the highest priority.

---

4. Using heapq.nsmallest() for Bottom-N Items

Similarly, heapq.nsmallest() allows you to find the n smallest items from a collection.

Copy

import heapq

# Sample data
data = [(1, 'task 1'), (3, 'task 3'), (2, 'task 2')]

# Get the two smallest items by priority
bottom_tasks = heapq.nsmallest(2, data, key=lambda x: x[0])

for task in bottom_tasks:
    print(f"Task: {task[1]}, Priority: {task[0]}")

Output:

Copy

Task: task 1, Priority: 1
Task: task 2, Priority: 2

This example retrieves the bottom 2 tasks with the lowest priority.

---

5. Using heapq.merge() for Merging Multiple Heaps

If you have multiple heaps and want to merge them into a single sorted iterator, you can use heapq.merge().

Copy

import heapq

# Two sorted lists (min-heaps)
heap1 = [1, 3, 5]
heap2 = [2, 4, 6]

# Merge the two heaps
merged = heapq.merge(heap1, heap2)

for item in merged:
    print(item)

Output:

Copy

1
2
3
4
5
6

The merge() function merges multiple sorted inputs (min-heaps) into a single sorted output.

---

6. Heap as a Min-Heap

The default behavior of heapq is to create a min-heap. Here’s an example of how you can use heapq to maintain a min-heap.

Copy

import heapq

# Create an empty list to represent the heap
heap = []

# Push elements into the heap
heapq.heappush(heap, 3)
heapq.heappush(heap, 1)
heapq.heappush(heap, 2)

# Pop elements from the heap (will return the smallest element first)
print(heapq.heappop(heap))  # Output: 1
print(heapq.heappop(heap))  # Output: 2
print(heapq.heappop(heap))  # Output: 3

Output:

Copy

1
2
3

In this example, heapq ensures that the smallest element is always at the root of the heap.

---

7. Heap for Sorting Large Data

You can use heapq to efficiently sort large data that might not fit entirely in memory at once.

Copy

import heapq

# Simulate a large dataset by using a generator
def large_data():
    yield from range(1000, 0, -1)

# Create a heap from the data
data = large_data()
heap = list(data)

# Convert the list into a heap (this will arrange the data as a min-heap)
heapq.heapify(heap)

# Pop the smallest element
while heap:
    print(heapq.heappop(heap), end=" ")

This example shows how you can take a large dataset and use a heap to efficiently sort and pop the smallest elements.

---

8. Heap for Priority Scheduling

You can implement a simple priority scheduling system using a heap.

Copy

import heapq

# Define tasks with their priority
tasks = []
heapq.heappush(tasks, (2, 'Task 2'))  # Priority 2
heapq.heappush(tasks, (1, 'Task 1'))  # Priority 1
heapq.heappush(tasks, (3, 'Task 3'))  # Priority 3

# Process tasks based on priority
while tasks:
    priority, task = heapq.heappop(tasks)
    print(f"Processing {task} with priority {priority}")

Output:

Copy

Processing Task 1 with priority 1
Processing Task 2 with priority 2
Processing Task 3 with priority 3

---

9. Finding Kth Largest Element

You can use heapq.nlargest() to find the Kth largest element in a list.

Copy

import heapq

data = [20, 15, 30, 10, 5, 25]

# Find the 3rd largest element using nlargest
third_largest = heapq.nlargest(3, data)[-1]

print(f"The 3rd largest element is: {third_largest}")

Output:

Copy

The 3rd largest element is: 20

---

10. Heap for Dijkstra’s Shortest Path Algorithm

A priority queue is often used in algorithms like Dijkstra’s to efficiently get the smallest edge weight.

Copy

import heapq

def dijkstra(graph, start):
    # The heap queue to store (distance, node) pairs
    pq = [(0, start)]  # distance to start node is 0
    distances = {start: 0}

    while pq:
        (dist, current_node) = heapq.heappop(pq)

        for neighbor, weight in graph[current_node]:
            distance = dist + weight

            # If the distance to the neighbor is shorter, update it
            if neighbor not in distances or distance < distances[neighbor]:
                distances[neighbor] = distance
                heapq.heappush(pq, (distance, neighbor))

    return distances

# Example graph (node: [(neighbor, weight), ...])
graph = {
    'A': [('B', 1), ('C', 4)],
    'B': [('C', 2), ('D', 5)],
    'C': [('D', 1)],
    'D': []
}

print(dijkstra(graph, 'A'))

Output:

Copy

{'A': 0, 'B': 1, 'C': 3, 'D': 4}

---