目录
问题
如何在一个集合中找到最大或最小的 N 个元素?
解决方案
使用 heapq
模块。
pip install heapq
heapq
模块中,有 nlargest()
以及 nsmallest()
两个函数:
import heapq
nums = [1, 8, 23, 2, 7, -4, 8, 18, 42, 37]
print(heapq.nlargest(3, nums))
print(heapq.nsmallest(3, nums))
# nlargest函数结构
nlargest(n, arr, key) # 其中n为取出的数量,arr为数组
# nsmallest函数结构
nsmallest(n, arr, key)
而由于这两个函数 nlargest()
nsmallest
都接受一个参数 key
,有了这个参数,就可以允许它们工作在更加复杂的数据结构之中。
import heapq
profolio = [
{"name":"IBM", "shares": 100, "price": 91.1},
{"name":"AAPL", "shares": 50, "price": 543.22},
{"name":"FB", "shares": 200, "price": 21.09},
{"name":"HPQ", "shares": 35, "price": 31.75},
{"name":"ACME", "shares": 75, "price": 115.65}
]
cheap = heapq.nsmallest(3, profolio, key=lambda s:s['price'])
print(cheap)
讨论
如果寻找集合 A A A 中最大最小的 N N N 个元素,且 N < < l e n ( A ) N<<len(A) N<<len(A),那么使用下述方案可以提供更好的性能。
首先函数会在底层将集合转化为列表,元素会以堆的顺序排列。
import heapq
nums = [1, 8, 2, 23, 7, -4, 18, 23, 42, 37, 2]
heap = list(nums)
heapq.heapify(heap)
print(heap)
上述代码 print(heap)
的结果为:
[-4, 2, 1, 23, 7, 2, 18, 23, 42, 37, 8]
在堆数据结构中,父节点和子节点的关系是基于数组的索引来确定的。对于一个给定的节点,其索引为 i i i,它的父节点、左子节点和右子节点的索引可以通过特定的公式计算得出。
而在 Python 的 heapq
模块实现的最小堆中,堆是一个列表,且堆属性满足对于所有索引
i
i
i(除了根节点,其索引为0),都有
h
e
a
p
[
i
]
>
=
h
e
a
p
[
(
i
−
1
)
/
/
2
]
heap[i] >= heap[(i-1)//2]
heap[i]>=heap[(i−1)//2]
即任何父节点的值都小于或等于其子节点的值。
[-4, 2, 1, 23, 7, 2, 18, 23, 42, 37, 8]
在这个堆中,
- 索引 0 的元素是 -4,它是根节点,没有父节点。
- 索引 1 的元素是 2,它的父节点是 -4(索引为
(1-1)//2 = 0
)。 - 索引 2 的元素是 1,它的父节点也是 -4(索引为
(2-1)//2 = 0
)。 - 索引 3 的元素是 23,它的父节点是 2(索引为
(3-1)//2 = 1
)。 - 索引 4 的元素是 7,它的父节点是 2(索引为
(4-1)//2 = 1
)。 - 索引 5 的元素是 2,它的父节点是 1(索引为
(5-1)//2 = 2
)。 - 以此类推,可以找到每个节点的父节点。
同样地,可以找到每个节点的子节点:
- 索引 0 的 -4 的左子节点是 2(索引为
2*0 + 1 = 1
),右子节点是 1(索引为2*0 + 2 = 2
)。 - 索引 1 的 2 的左子节点是 23(索引为
2*1 + 1 = 3
),右子节点是 7(索引为2*1 + 2 = 4
)。 - 索引 2 的 1 的左子节点是 2(索引为
2*2 + 1 = 5
),但没有右子节点(因为索引2*2 + 2 = 6
处没有元素)。 - 以此类推,可以找到每个节点的子节点。
Python 的 heapq 模块提供了一个 heapify
函数,该函数能够将一个可变的列表转换为最小堆。过程是自动的,但是可以进行模拟:
- 首先找到最后一个非叶子节点的索引:
- 在一个完全二叉树中,最后一个非叶子节点的索引是
len(heap) // 2 - 1
- 而本列表中,
len(nums) = 11
,所以最后一个非叶子节点的索引是11 // 2 - 1 = 4
- 在一个完全二叉树中,最后一个非叶子节点的索引是
- 从最后一个非叶子节点开始,向上进行堆化(heapify):
- 对于每个节点,比较它与它的子节点的值。
- 如果节点的值大于其子节点中的最小值,则交换这个节点与其最小子节点的值。
- 重复这个过程,直到堆的根节点。
- 重复步骤2,直到整个列表满足堆属性。
手动模拟过程:
# 初始列表
heap = [1, 8, 2, 23, 7, -4, 18, 23, 42, 37, 2]
# 从最后一个非叶子节点开始,向上进行堆化
last_non_leaf = len(heap) // 2 - 1
for i in range(last_non_leaf, -1, -1):
# 比较当前节点与其子节点的值,并进行交换
current = heap[i]
left_child_idx = 2 * i + 1
right_child_idx = 2 * i + 2
smallest = i
# 如果左子节点存在且小于当前节点,更新最小值索引
if left_child_idx < len(heap) and heap[left_child_idx] < current:
smallest = left_child_idx
# 如果右子节点存在且小于最小值,更新最小值索引
if right_child_idx < len(heap) and heap[right_child_idx] < heap[smallest]:
smallest = right_child_idx
# 如果当前节点不是最小值,交换它们
if smallest != i:
heap[i], heap[smallest] = heap[smallest], heap[i]
# 最终的堆
print(heap)
执行上述代码后,我们得到了一个最小堆。这个过程是迭代的,从最后一个非叶子节点开始,向上逐步将每个节点与其子节点进行比较和交换,直到整个列表满足堆属性。
标签:heapq,S01E03,idx,23,Python,索引,heap,Cookbook,节点 From: https://blog.csdn.net/weixin_43098506/article/details/139298770