【Python Cookbook】S01E03 找到最大最小的N个元素

news2024/10/4 18:21:16

问题

如何在一个集合中找到最大或最小的 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$ 中最大最小的 $N$ 个元素，且 $N << l e n (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$ ，它的父节点、左子节点和右子节点的索引可以通过特定的公式计算得出。

而在 Python 的 heapq 模块实现的最小堆中，堆是一个列表，且堆属性满足对于所有索引 $i$ （除了根节点，其索引为0），都有 $h e a p [i] >= h e a p [(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)