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

需求:
快速的获取一个列表中最大/最小的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官方文档

原文地址:https://www.cnblogs.com/valorchang/p/11289115.html