要想找出一组值中的最大、最小的 n
个值,就要利用到 heapq
中的 nlargest()
和 nsmallest()
两个函数:
1 2 3 4 5
| import heapq nums = [1, 9, 6, 8, 4, 2, 5, 7, 0] print(heapq.nlargest(3, nums)) print(heapq.nsmallest(3, nums))
|
这两个函数都可以接受一个参数 key
,从而允许它们工作在更加复杂的数据结构之上:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21
| 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']) print(cheap) print(expensive)
|
上面的代码似乎有些难于理解,从 StackOverflow 了解到,lambada
的作用是创建一个匿名函数,lambda
的后面跟着参数名列表,然后再在其后跟着一个代码块,参数名列表和代码块之间用分号隔开,这样子的结构在 Python
中与 while
、for
和 if
类似,它们都是拥有代码块的典型代表,lambda
只是另外一个代表而已。我们可以比较一下利用 lambda
和 def
定义函数的不同:
1 2 3 4 5
| adder_lambda = lambda parameter1, parameter2: parameter1 + parameter2 def adder_regular(parameter1, parameter2): return parameter1 + parameter2
|
lambda
给我们提供了一种完成类似这样子任务的方式而不用为函数取一个名字,相对于一个 def
定义的函数来说这样子要好得多。
对于 nlargest
和 nsmallest
两个函数中增加 key
参数,根据 Python documentation
中描述的那样,上面的两条语句分别相当于:
1 2
| cheap = sorted(portfolio, key=lambda s: s['price'])[:3] expensive = sorted(portfolio, key=lambda s: s['price'], reverse=True)[:3]
|
为了验证正确性,不妨运行一下下面的代码:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19
| 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']) cheap_copy = sorted(portfolio, key=lambda s: s['price'])[:3] expensive = heapq.nlargest(3, portfolio, key=lambda s: s['price']) expensive_copy = sorted(portfolio, key=lambda s: s['price'], reverse=True)[:3] print(cheap == cheap_copy) print(expensive == expensive_copy)
|
由此可见上面的说法是完全正确的。
需要特别注意的是,在 nlargest
和 nsmallest
函数中以及在 sorted
中的 key
返回的必须是具体的数值,切莫想当然地以为是 price
,而是 price
的 value
。
堆最重要的特性就是 heap[0]
总是最小的那个元素,此外,接下来的最小元素可以依次通过 heapq.pop()
方法轻松找到,该方法会将第一个元素(最小的)弹出,然后会以第二小的元素取而代之。当所要找的元素数量 N 相对较小时,函数 nlargest()
和 nsmallest()
才是最适用的。如果只是简单地想找到最小或最大的元素(N=1时),那么用 min()
和 max()
会更加快。同样,如果 N 和集合本身的大小差不多大,通常更快的方法是先对集合排序,然后做切片操作(例如使用 sorted(items)[:N]
或者 sorted(items)[:-N]
。应该要注意的是,nlargest()
和 nsmallest()
的实际实现会根据使用它们的方式而有所不同,可能会相应作出一些优化措施(比如,当 N 的大小同输入大小很接近时,就会采用排序方法)。
可以使用 heapify
将一个列表转换成堆。
1 2 3 4 5
| import heapq nums = [1, 8, 2, 23, 7, -4, 18, 23, 42, 37, 2] heap = list(nums) heapq.heapify(heap) print(heap)
|