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

[fanqieipnet]

我们在写查询元素的代码时，通常会使用包含 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 获取前几个最小值。