要想找出一组值中的最大、最小的 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)) # [9, 8, 7]
print(heapq.nsmallest(3, nums)) # [0, 1, 2]

这两个函数都可以接受一个参数 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)

# 结果为:
# [{'name': 'YHOO', 'shares': 45, 'price': 16.35}, {'name': 'FB', 'shares': 200, 'price': 21.09}, {'name': 'HPQ', 'shares': 35, 'price': 31.75}]
# [{'name': 'AAPL', 'shares': 50, 'price': 543.22}, {'name': 'ACME', 'shares': 75, 'price': 115.65}, {'name': 'IBM', 'shares': 100, 'price': 91.1}]

上面的代码似乎有些难于理解,从 StackOverflow 了解到,lambada 的作用是创建一个匿名函数,lambda 的后面跟着参数名列表,然后再在其后跟着一个代码块,参数名列表和代码块之间用分号隔开,这样子的结构在 Python 中与 whileforif类似,它们都是拥有代码块的典型代表,lambda 只是另外一个代表而已。我们可以比较一下利用 lambdadef 定义函数的不同:

1
2
3
4
5
adder_lambda = lambda parameter1, parameter2: parameter1 + parameter2


def adder_regular(parameter1, parameter2):
return parameter1 + parameter2

lambda 给我们提供了一种完成类似这样子任务的方式而不用为函数取一个名字,相对于一个 def 定义的函数来说这样子要好得多。

对于 nlargestnsmallest 两个函数中增加 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) # True
print(expensive == expensive_copy) # True

由此可见上面的说法是完全正确的。

需要特别注意的是,在 nlargestnsmallest 函数中以及在 sorted 中的 key 返回的必须是具体的数值,切莫想当然地以为是 price,而是 pricevalue

堆最重要的特性就是 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) # [-4, 2, 1, 23, 7, 2, 18, 23, 42, 37, 8]