Python找到最大或最小的N个元素

问题

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

解决方案

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

import heapq
nums = [2, 6, 1, -3, 56, 37, 18, 9, 38, 42]
print(heapq.nlargest(3, nums))    # prints 3 largest nums: [56, 42, 38]
print(heapq.nsmallest(3, nums))   # prints 3 smallest nums: [-3, 1, 2]

这两个函数原型如下：

def nlargest(n, iterable, key=None):
    """Find the n largest elements in a dataset.
    Equivalent to:  sorted(iterable, key=key, reverse=True)[:n]
    """
def nsmallest(n, iterable, key=None):
    """Find the n smallest elements in a dataset.
    Equivalent to:  sorted(iterable, key=key)[:n]
    """

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

import heapq
portfolio = [
{'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': 'YHOO', 'shares': 45, 'price': 16.35},
{'name': 'ACME', 'shares': 75, 'price': 115.65}
]
cheap = heapq.nsmallest(3, portfolio, key=lambda s: s['price'])
expensive = heapq.nlargest(3, portfolio, key=lambda s: s['price'])

讨论

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

>>> import heapq
>>> nums = (2, 6, 1, -3, 56, 37, 18, 9, 38, 42)
>>> heap = list(nums)
>>> heapq.heapify(heap)
>>> heap
[-3, 2, 1, 6, 42, 37, 18, 9, 38, 56]
>>>

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

>>> heapq.heappop(heap)
-3
>>> heapq.heappop(heap)
1
>>> heapq.heappop(heap)
2
>>>

当所要找的的元素数量相对较小时，函数nsmallest()和nlargest()才是最适用的。如果只是简单的想找到最小或最大的元素（N=1）时，那么用min()和max()会更加快。同样，如果N和集合本身的大小差不多大，通常更快的方法是先对数据集排序，然后做切片操作（例如，使用sorted(items)[:N]或者sorted(items)[-N:]）。应该要注意的是，nsmallest()和nlargest()的实际实现会根据使用它们的方式而有所不同，可能会相应作出一些优化措施（比如，当N的大小同输入大小很接近时，就会采用排序的方法）。

本文并没有讨论堆数据结构的底层实现细节，通常在优秀的算法和数据结构相关的书籍里都能找到对数据结构的实现方法。另外，如需更详细的了解使用heapq模块，参考官方文档：https://docs.python.org/3.0/library/heapq.html