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