【Python 语法】heapq 模块

heapq 是 Python 标准库中用于实现堆队列（heap queue）算法的模块。堆队列是一个基于堆（heap）数据结构的优先队列，它能在 O(log n) 时间内执行插入、删除最小元素等操作。

Python 中的 heapq 模块实现的是一个最小堆（min-heap），即 堆顶元素是堆中的最小元素。

堆的应用场景

1. 优先队列：heapq 可以用来实现优先队列，按优先级顺序处理元素。
2. K 小/大元素问题：通过 heapq.nsmallest() 或 heapq.nlargest()，你可以快速找到数据中的最小或最大 K 个元素。
3. 合并多个排序的迭代器：使用 heapq.merge() 可以高效地合并多个已排序的迭代器。

主要功能

heapq 模块提供了以下几个常用的函数：

1. heapq.heappush(heap, item)：将一个元素 item 添加到堆 heap 中，保持堆的特性。

   import heapq
   heap = []
   heapq.heappush(heap, 5)
   heapq.heappush(heap, 3)
   heapq.heappush(heap, 8)
   print(heap)  # 输出: [3, 5, 8]
   

2. heapq.heappop(heap)：移除并返回堆中的最小元素（即堆顶元素），然后重新调整堆。

   import heapq
   heap = [3, 5, 8]
   min_item = heapq.heappop(heap)
   print(min_item)  # 输出: 3
   print(heap)      # 输出: [5, 8]
   

3. heapq.heappushpop(heap, item)：先将元素 item 添加到堆中，再返回并移除堆中的最小元素。这个操作比先执行 heappush() 再执行 heappop() 更高效。

   import heapq
   heap = [3, 5, 8]
   result = heapq.heappushpop(heap, 2)
   print(result)  # 输出: 2
   print(heap)    # 输出: [3, 5, 8]
   

4. heapq.heapreplace(heap, item)：先移除堆顶元素，再添加 item 到堆中，返回移除的堆顶元素。

   import heapq
   heap = [3, 5, 8]
   result = heapq.heapreplace(heap, 1)
   print(result)  # 输出: 3
   print(heap)    # 输出: [1, 5, 8]
   

5. heapq.nsmallest(n, iterable, key=None)：返回 iterable 中最小的 n 个元素，按升序排序。

   import heapq
   nums = [5, 1, 8, 2, 9, 3]
   result = heapq.nsmallest(3, nums)
   print(result)  # 输出: [1, 2, 3]
   

6. heapq.nlargest(n, iterable, key=None)：返回 iterable 中最大的 n 个元素，按降序排序。

   import heapq
   nums = [5, 1, 8, 2, 9, 3]
   result = heapq.nlargest(3, nums)
   print(result)  # 输出: [9, 8, 5]
   

示例：使用 heapq 实现优先队列

import heapq

# 创建一个空堆
priority_queue = []

# 向优先队列中添加任务，任务由 (优先级, 任务内容) 的元组表示
heapq.heappush(priority_queue, (1, "Task 1"))
heapq.heappush(priority_queue, (3, "Task 3"))
heapq.heappush(priority_queue, (2, "Task 2"))

# 按优先级顺序处理任务
while priority_queue:
    priority, task = heapq.heappop(priority_queue)
    print(f"Processing: {task} with priority {priority}")