如何使用内置模块heapq查找最大或最小的 N 个元素？

我们在写查询元素的代码时，通常会使用包含 yield 表达式的生成器函数，查找最大或最小的 N 个元素时，表面上这是一个很简单的话题，其实如果要考虑的全面，也是需要留意一些事情的。今天番茄加速就来讲一下。

当查找元素个数N = 1时，建议直接使用max或min方法

当查找元素个数接近整个列表长度时，建议使用sorted函数以切片的方式获取

当要查找的元素个数相对比较小的时候，函数 nlargest() 和 nsmallest() 是很合适的

相信大家都对前两种情况的解决方法比较熟悉，第三种使用内置模块heapq是算法中的堆结构，常见的大根堆，小根堆，

>>> nums = [1, 8, 2, 23, 7, -4, 18, 23, 42, 37, 2]

>>> import heapq

>>> heap = list(nums)

>>> heapq.heapify(heap)

>>> heap

[-4, 2, 1, 23, 7, 2, 18, 23, 42, 37, 8]

>>>

Python中heapify后，默认建立一个小根堆。它最重要的特征是 heap[0] 永远是最小的元素。

比如，如果想要查找最小的 3 个元素，你可以这样做，首先执行一次heappop后，次小元素变为最小，如下图所示：

>>> heapq.heappop(heap)

-4

再次执行两次后，就能得到列表的前三个最小的元素为[-4,1,2]：

>>> heapq.heappop(heap)

1

>>> heapq.heappop(heap)

2

当然，也可以直接使用nsmallest获取前几个最小值。