Python heapq模块

一、堆与 heapq 模块简介

heapq 是 Python 内置的堆队列算法模块，基于最小堆​（父节点 ≤ 子节点）实现优先队列，支持快速插入、删除和极值查询操作。

核心特点

• ​时间复杂度：
  • 插入/删除：O(log n)
  • 极值查询：O(1)
  • 堆化（heapify）：O(n)
• ​默认最小堆，需手动模拟最大堆（通过取负数）
• ​原地操作：直接修改列表，无需额外内存

---

​二、核心函数与用法

​1. 堆的创建与维护

函数/方法	功能描述	示例
heapq.heapify(iterable)	将列表原地转换为堆结构（满足堆属性）	data = [3,1,4]; heapq.heapify(data) → [1,3,4]
heapq.heappush(heap, item)	插入元素并维护堆结构	heap = []; heapq.heappush(heap, 5) → 堆变为 [5]
heapq.heappop(heap)	弹出并返回堆顶（最小元素）	val = heapq.heappop(heap) → val=1，堆变为 [3,4]
heapq.heappushpop(heap, item)	先插入再弹出最小元素（比分开调用高效）	heapq.heappushpop([2,5], 3) → 返回 2，堆变为 [3,5]
heapq.heapreplace(heap, item)	先弹出最小元素再插入新元素（堆不能为空）	heapq.heapreplace([2,5], 3) → 返回 2，堆变为 [3,5]

​2. 极值查询与合并

函数/方法	功能描述	示例
heapq.nlargest(n, iterable)	返回前 n 个最大元素（降序）	nlargest(3, [1,5,3]) → [5,3,1]
heapq.nsmallest(n, iterable)	返回前 n 个最小元素（升序）	nsmallest(2, [5,1,3]) → [1,3]
heapq.merge(*iterables)	合并多个已排序序列（返回迭代器）	list(merge([1,3], [2,4])) → [1,2,3,4]

---

​三、高级用法与场景

​1. 实现最大堆

通过存储负数模拟最大堆：

import heapq

data = [3, 1, 4]
max_heap = [ -x for x in data ]  # 存储负数
heapq.heapify(max_heap)

# 插入新元素
heapq.heappush(max_heap, -10)

# 弹出最大元素
max_val = -heapq.heappop(max_heap)  # 输出 10

2. 自定义优先级排序​​

使用元组 (优先级, 数据) 控制排序，适用于任务调度等场景：

tasks = []
heapq.heappush(tasks, (2, '低优先级任务'))
heapq.heappush(tasks, (1, '高优先级任务'))

# 弹出最高优先级任务
priority, task = heapq.heappop(tasks)  # task='高优先级任务'

​3. 高效解决 Top K 问题

def top_k_elements(arr, k):
    heap = []
    for num in arr:
        heapq.heappush(heap, num)
        if len(heap) > k:
            heapq.heappop(heap)
    return heap  # 保留最大的 k 个元素

对比传统排序方法，时间复杂度从 O(n log n) 优化为 O(n log k)。

---

​四、典型应用场景

​1. 优先级队列

实现任务调度系统：

class PriorityQueue:
    def __init__(self):
        self._queue = []
        self._index = 0  # 处理同优先级任务的顺序

    def push(self, item, priority):
        heapq.heappush(self._queue, (-priority, self._index, item))
        self._index += 1

    def pop(self):
        return heapq.heappop(self._queue)[-1]

高优先级任务（数值越小优先级越高）优先出队。

​2. 数据流实时处理

实时获取最新数据中的最小元素：

stream = [5, 3, 9, 1, 7]  # 模拟实时数据流
heap = []
for num in stream:
    heapq.heappush(heap, num)
    print("当前最小元素:", heap[0])  # O(1) 查询

​3. 算法优化

• ​Dijkstra 最短路径：用堆优化节点选择步骤，时间复杂度从 O(V^2) 降至 O(E + V log V)。
• ​合并 K 个有序链表：利用 heapq.merge 高效合并。

---

​五、注意事项

1. ​列表必须满足堆属性：直接修改列表可能导致堆结构破坏，需通过 heappush/heappop 维护。
2. ​最大堆的实现限制：需手动处理负数，复杂对象需自定义键值。
3. ​性能权衡：
   • nlargest/nsmallest 适合少量极值查询（如 Top 10）
   • 频繁动态插入/删除时，优先使用堆而非全排序

---

​六、完整示例

​堆排序实现

def heap_sort(arr):
    heapq.heapify(arr)
    return [heapq.heappop(arr) for _ in range(len(arr))]

data = [3, 1, 4, 1, 5]
sorted_data = heap_sort(data)  # [1, 1, 3, 4, 5]

​合并多个有序序列

a = [1, 3, 5]
b = [2, 4, 6]
merged = list(heapq.merge(a, b))  # [1, 2, 3, 4, 5, 6]

---

掌握 heapq 可显著提升数据处理效率，尤其适用于动态数据管理和极值查询场景！