Python中heapq
库的基础使用方法和示例代码,包含详细注释说明:
1. 基本功能
heapq
实现的是最小堆(父节点值 ≤ 子节点值),核心操作包括:
- 插入元素 :
heappush(heap, item)
- 弹出最小值 :
heappop(heap)
- 堆化列表 :
heapify(list)
(将无序列表转换为堆) - 查看最小值 :
heap[0]
2. 基础示例代码
python
import heapq
# 创建一个空堆
heap = []
# 插入元素
heapq.heappush(heap, 3)
heapq.heappush(heap, 1)
heapq.heappush(heap, 4)
print("当前堆:", heap) # 输出: [1, 3, 4](实际存储为堆结构)
# 弹出最小值
min_val = heapq.heappop(heap)
print("弹出的最小值:", min_val) # 输出: 1
print("剩余堆:", heap) # 输出: [3, 4]
# 堆化现有列表
existing_list = [5, 2, 7, 1]
heapq.heapify(existing_list)
print("堆化后的列表:", existing_list) # 输出: [1, 2, 7, 5]
3. 进阶用法
3.1 合并堆
python
heap1 = [1, 3, 5]
heap2 = [2, 4, 6]
merged_heap = heapq.merge(heap1, heap2)
print("合并后的堆(迭代器):", list(merged_heap)) # 输出: [1, 2, 3, 4, 5, 6]
3.2 处理复杂元素
当元素是元组时,默认按第一个元素排序。可通过tuple
调整优先级:
python
# 按元组第二个元素排序(需预处理)
tasks = [(3, "任务A"), (1, "任务B"), (2, "任务C")]
heapq.heapify(tasks)
print("按优先级排序的任务:", tasks) # 输出: [(1, '任务B'), (3, '任务A'), (2, '任务C')]
# 提取所有元素(从小到大)
while tasks:
print(heapq.heappop(tasks))
# 输出:
# (1, '任务B')
# (2, '任务C')
# (3, '任务A')
3.3 限制堆大小
保留堆中最大的N个元素(或最小的N个元素):
python
# 保留最小的3个元素
nums = [4, 1, 7, 3, 8, 5]
heapq.heapify(nums)
heapq.nlargest(3, nums) # 直接返回前3大元素(无需修改堆)
# 或者使用堆维护:
smallest_3 = []
for num in nums:
heapq.heappush(smallest_3, num)
if len(smallest_3) > 3:
heapq.heappop(smallest_3)
print("最小的3个元素:", smallest_3) # 输出: [1, 3, 4]
4. 注意事项
- 堆的索引 :堆的根节点在索引0,子节点在
2*i+1
和2*i+2
。 - 时间复杂度 :
- 插入/弹出:
O(log n)
- 堆化列表:
O(n)
- 插入/弹出:
- 直接操作堆 :避免手动修改堆列表,需通过
heapq
函数维护堆性质。
5. 应用场景
- 优先队列:如任务调度系统。
- Top K 问题:快速找到前K大/小的元素。
- 图算法:如Dijkstra算法中的优先队列。
- 流数据处理:实时维护最大/最小值的集合。
如果需要最大堆,可以将元素取负数后存入最小堆,使用时再转换回来。