126. Complex Data Structures with heapq

The heapq module in Python provides a way to implement priority queues using heaps. A heap is a specialized tree-based data structure that satisfies the heap property, where the parent node is smaller (in a min-heap) or larger (in a max-heap) than its children. The heapq module specifically implements a min-heap by default, but max-heaps can be simulated.

Here are 10 Python code snippets demonstrating the use of heapq to manage complex data structures like priority queues.

1. Basic Min-Heap Usage

Copy

import heapq

heap = []
heapq.heappush(heap, 20)
heapq.heappush(heap, 5)
heapq.heappush(heap, 15)

print(heap)  # Output: [5, 20, 15]
print(heapq.heappop(heap))  # Output: 5
print(heap)  # Output: [15, 20]

---

2. Max-Heap Simulation with Negatives

Copy

import heapq

max_heap = []
heapq.heappush(max_heap, -20)
heapq.heappush(max_heap, -5)
heapq.heappush(max_heap, -15)

print([-item for item in max_heap])  # Output: [20, 5, 15]
print(-heapq.heappop(max_heap))  # Output: 20
print([-item for item in max_heap])  # Output: [15, 5]

---

3. Priority Queue with Custom Objects

Copy

import heapq

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

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

task1 = Task(2, 'Task 2')
task2 = Task(1, 'Task 1')
task3 = Task(3, 'Task 3')

tasks = []
heapq.heappush(tasks, task1)
heapq.heappush(tasks, task2)
heapq.heappush(tasks, task3)

while tasks:
    task = heapq.heappop(tasks)
    print(f"Task: {task.task_name}, Priority: {task.priority}")

---

4. Heapq for Sorting

Copy

import heapq

numbers = [7, 9, 3, 6, 2, 8, 4]
heapq.heapify(numbers)  # Turn the list into a heap

sorted_numbers = [heapq.heappop(numbers) for _ in range(len(numbers))]
print(sorted_numbers)  # Output: [2, 3, 4, 6, 7, 8, 9]

---

5. Heapq with heapreplace to Replace Minimum Element

Copy

import heapq

numbers = [7, 9, 3, 6, 2, 8, 4]
heapq.heapify(numbers)

# Replace the smallest element (2) with 5
heapq.heapreplace(numbers, 5)
print(numbers)  # Output: [3, 5, 4, 6, 9, 8, 7]

---

6. Efficient Merging of Multiple Sorted Lists

Copy

import heapq

list1 = [1, 3, 5, 7]
list2 = [2, 4, 6, 8]
list3 = [0, 9, 10]

merged_list = heapq.merge(list1, list2, list3)
print(list(merged_list))  # Output: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

---

7. Find the N Largest Elements in a List

Copy

import heapq

numbers = [1, 3, 5, 7, 9, 2, 4, 6, 8]
n_largest = heapq.nlargest(3, numbers)
print(n_largest)  # Output: [9, 8, 7]

---

8. Find the N Smallest Elements in a List

Copy

import heapq

numbers = [1, 3, 5, 7, 9, 2, 4, 6, 8]
n_smallest = heapq.nsmallest(3, numbers)
print(n_smallest)  # Output: [1, 2, 3]

---

9. Custom Priority Queue with Tuple and Custom Comparator

Copy

import heapq

# A custom comparator that sorts by the second element of the tuple
data = [(3, "task3"), (1, "task1"), (2, "task2")]

# Min-heap based on the first element of the tuple
heapq.heapify(data)

while data:
    task = heapq.heappop(data)
    print(task)  # Output: (1, 'task1'), (2, 'task2'), (3, 'task3')

---

10. Kth Smallest Element in a Stream

Copy

import heapq

class KthSmallest:
    def __init__(self, k, nums):
        self.k = k
        self.min_heap = nums
        heapq.heapify(self.min_heap)

        while len(self.min_heap) > k:
            heapq.heappop(self.min_heap)

    def add(self, val):
        heapq.heappush(self.min_heap, val)
        if len(self.min_heap) > self.k:
            heapq.heappop(self.min_heap)
        return self.min_heap[0]

kth_smallest = KthSmallest(3, [4, 5, 8, 2])
print(kth_smallest.add(3))  # Output: 4
print(kth_smallest.add(5))  # Output: 5
print(kth_smallest.add(10)) # Output: 5

---

These examples demonstrate various ways to use the heapq module for implementing efficient priority queues, merging sorted data, finding largest or smallest elements, and more complex applications with custom objects. The heapq module is especially useful when you need to manage a dynamic collection of elements and perform fast retrieval of the smallest or largest elements in the collection.