heapq
是 Python 标准库中的一个模块,它提供了堆队列算法的实现,也称为优先队列算法。
一、非常值得注意的是:heapq在Python中默认实现的是小根堆(也称为最小堆)
堆队列是一种树形数据结构,可以用数组或类似数组的对象(如 Python 列表)来表示。heapq
模块提供了构建和操作最小堆(min-heap)的功能,但也可以用于实现最大堆(max-heap)。
二、heapq
模块中的函数主要有以下几个:
heapq.heappush(heap, item)
: 将元素item
添加到堆heap
中,保持堆的不变性。时间复杂度为 O(log n)。heapq.heappop(heap)
: 弹出堆heap
中的最小元素,并返回它。如果堆为空,则引发IndexError
。时间复杂度为 O(log n)。heapq.heappushpop(heap, item)
: 先将元素item
推入堆heap
,然后弹出并返回堆中的最小元素。这两个操作组合起来的时间复杂度为 O(log n),这比先调用heappush()
再调用heappop()
更快。heapq.heapify(x)
: 将列表x
转换为一个堆,使其满足堆的性质。也就是说,x[0]
是堆中的最小元素。这个函数假设列表x
是可索引的,并且仅在其上执行原地操作,不会生成新的列表。heapq.heapreplace(heap, item)
: 弹出并返回堆heap
中的最小元素,同时将新元素item
推入堆中。这相当于heappop(heap)
后立即调用heappush(heap, item)
,但更为高效,因为它可以只用一次树旋转来完成操作。heapq.nlargest(n, iterable, key=None)
: 返回可迭代对象iterable
中最大的n
个元素,作为一个列表返回。如果n
大于或等于iterable
的长度,则返回iterable
的所有元素,并按降序排列。key
参数指定一个单参数函数,用于从iterable
的每个元素中提取比较键(例如,key=str.lower
)。默认值为None
(直接比较元素)。
这些函数使得 heapq
模块成为实现诸如堆排序、Dijkstra 的最短路径算法和 Prim 的最小生成树算法等算法的有力工具。同时,由于堆是一种优先队列,因此 heapq
模块也常用于需要高效处理具有优先级的数据的场景中。
三、要将heapq设置为大根堆(也称为最大堆),即父节点的值总是大于或等于其子节点的值,可以通过以下方式实现:
- 存储负数:
- 当你向堆中添加元素时,可以存储这些元素的相反数(即负数)。这样,由于heapq是小根堆,存储的负数越大(即原数越小),其位置就越靠后。
- 当你需要从堆中获取元素时,同样取出的是负数,此时需要取这个负数的相反数,即得到了实际的最大值。
示例代码:
import heapq
# 原始数据
data = [1, 5, 3, 2, 9, 5]
# 转换为负数并构建小根堆
neg_data = [(-i, i) for i in data] # 存储负数和原始数据的元组
heapq.heapify(neg_data) # 将neg_data列表根据第一个元素转化为堆
# 弹出并返回最大元素(即负数最小的元素的原始值)
while neg_data:
max_val = heapq.heappop(neg_data)
print(max_val) # 输出(-最大值,最大值)