Python Cookbook 之一 数据结构和算法（四）： 找到最大或最小的 N 个元素

欢迎大家扫码关注我的微信公众号：

一、问题

我们想在某个集合中找出最大或最小的 N 个元素。

二、解决方案

heapq 模块中有两个函数 —— nlargest() 和 nsmallest() —— 它们正是我们所需要的：

>>> import heapq
>>> lst = [1, 22, -23, -232, 789, 230, 2907]
>>> heapq.nlargest(2, lst)
[2907, 789]
>>> heapq.nsmallest(2, lst)
[-232, -23]

这两个函数都可以接受一个参数 key， 从而允许它们工作在更加复杂的数据结构之上：

>>> import heapq
>>> lst = [{'name': 'CookBook', 'price': 108, 'balance': 11}, {'name': 'Mysql', 'price': 78, 'balance': 232}, {'name': 'JavaScript', 'price': 22, 'balance': 565}]
>>> cheapest = heapq.nsmallest(2, lst, key=lambda x: x['price'])
>>> cheapest
[{'name': 'JavaScript', 'price': 22, 'balance': 565}, {'name': 'Mysql', 'price': 78, 'balance': 232}]
>>> expensivest = heapq.nlargest(2, lst, key=lambda x: x['price'])
>>> expensivest
[{'name': 'CookBook', 'price': 108, 'balance': 11}, {'name': 'Mysql', 'price': 78, 'balance': 232}]

三、讨论

如果正在寻找最大或最小的 N 个元素， 且同集合中元素的总数目相比， N 很小， 那么下面这些函数可以提供更好的性能。 这些函数首先会在底层将数据转化成列表， 且元素会以堆的顺序排列：

>>> lst = [1, 22, -23, -232, 789, 230, 2907]
>>> import heapq
>>> tmp = list(lst)
>>> heapq.heapify(tmp)
>>> tmp
[-232, 1, -23, 22, 789, 230, 2907]

堆最重要的特性就是 heap[0] 总是最小的那个元素。 此外， 接下来的元素可依次通过 heapq.heappop() 方法轻松找到。 该方法会将第一个元素（最小的）弹出， 然后第二小的元素取而代之（这个操作的复杂度是 O(log N)， N 代表堆的大小）。

>>> heapq.heappop(tmp)  # 想找最小的元素
-232
>>> heapq.heappop(tmp)  # 想找第二小的元素
-23
>>> heapq.heappop(tmp)  # 想找第三小的元素
1

【注意】当我们想要找的最大或最小的 N 个元素比这个集合中所有元素的个数要小很多的时候， 我们利用上述的方法， 会得到更好的性能。

【注意】

1. 当所要找的元素数量相对较小时， 函数 nlargest() 和 nsmallest() 才是最适用的。
2. 如果只是简单的想找到最大或最小的那一个元素时， 用 min() 和 max() 会更快。
3. 同样如果 N 和集合本身的大小差不多时， 通常更快的方法是先对集合排序， 然后做切片操作。

应该注意的是， nlargest() 和 nsmallest() 的实际实现会根据使用它们的方式不同而不同， 可能会相应作出一些优化措施。