Python实用黑科技——找出最大/最小的n个元素

Posted 2022-06-05 valorchang

需求：
快速的获取一个列表中最大/最小的n个元素。

方法：
最简便的方法是使用heapq模组的两个方法nlargest()和nsmallest()，例如：

In [1]: import heapq
In [2]: nums = [1, 0, -23, 45, 34, -11, 0, 2, 99, 103, -78]
In [3]: print(heapq.nlargest(3, nums))
[103, 99, 45]
In [4]: print(heapq.nsmallest(3, nums))
[-78, -23, -11]
这两个方法还可以根据具体的key值来获取复杂数据结构的结果， 例如：

In [5]: computers = [
...: ‘name‘ : ‘IBM‘, ‘shares‘ : 100, ‘price‘ : 91.1,
...: ‘name‘ : ‘APPL‘, ‘shares‘ : 59, ‘price‘ : 543.22,
...: ‘name‘ : ‘FB‘, ‘shares‘ : 200, ‘price‘ : 21.09,
...: ‘name‘ : ‘HPQ‘, ‘shares‘ : 35, ‘price‘ : 31.75,
...: ‘name‘ : ‘YHOO‘, ‘shares‘ : 45, ‘price‘ : 16.35,
...: ‘name‘ : ‘ACME‘, ‘shares‘ : 75, ‘price‘ : 115.65
...: ]
In [6]: cheap = heapq.nsmallest(4, computers, key=lambda s:s[‘price‘])
In [7]: expensive = heapq.nlargest(4, computers, key=lambda s:s[‘price‘])

In [8]: cheap
Out[8]:
[‘name‘: ‘YHOO‘, ‘price‘: 16.35, ‘shares‘: 45,
‘name‘: ‘FB‘, ‘price‘: 21.09, ‘shares‘: 200,
‘name‘: ‘HPQ‘, ‘price‘: 31.75, ‘shares‘: 35,
‘name‘: ‘IBM‘, ‘price‘: 91.1, ‘shares‘: 100]

In [9]: expensive
Out[9]:
[‘name‘: ‘APPL‘, ‘price‘: 543.22, ‘shares‘: 59,
‘name‘: ‘ACME‘, ‘price‘: 115.65, ‘shares‘: 75,
‘name‘: ‘IBM‘, ‘price‘: 91.1, ‘shares‘: 100,
‘name‘: ‘HPQ‘, ‘price‘: 31.75, ‘shares‘: 35]
扩展：
可以看出这两个方法对于找出n个最大/最小值的方便，从本质上看，其实调用模组heapq的左右就是返回一个排序堆构成的列表：

In [10]: nums = [1, 3, 5, -4, 18, -23, 0, 9, -6, 11]
In [11]: heap = list(nums)
In [12]: heapq.heapify(heap)
In [13]: heap
Out[13]: [-23, -6, 0, -4, 11, 5, 1, 9, 3, 18]

同样的，连续的元素也可以很方便的通过另一个方法heappop来获取：

In [14]: heapq.heappop(heap)
Out[14]: -23
In [15]: heapq.heappop(heap)
Out[15]: -6
In [16]: heapq.heappop(heap)
Out[16]: -4

当然，如果要选择最大/最小的1个元素(n=1,其实就是最大最小值），使用min(), max()自然更快。此外，对于获取n个元素，直接采用sorted(items)[:n]或者sorted(items)[-n:]也是方便快捷的。虽然替代方法不少，但是对于堆得使用确实很有意思的，有兴趣的读者可以继续研究它的底层实现。同样，对于heapq这个模组的使用可以参考heapq官方文档