从一个集合中查找最大最小的N个元素——Python heapq 堆数据结构-白红宇

Top N问题在搜索引擎、推荐系统领域应用很广，如果用我们较为常见的语言，如C、C++、Java等，代码量至少也得五行，但是用Python的话，只用一个函数就能搞定，只需引入heapq(堆队列)这个数据结构即可。今天偶然看到这个库，特意记下之。

先看一个例子：

1 >>> import heapq2 >>> nums = [1,8,2,23,7,-4,18,23,42,37,2]3 >>> print heapq.nlargest(3, nums)4 [42, 37, 23]5 >>> 6 >>> print heapq.nsmallest(3, nums)7 [-4, 1, 2]

是不是很简洁？

我们具体来看一下具体的函数定义。heapq有很多函数，最为堆，队列，可想而知，也就是那些push，pop之类的操作，详细请看官方文档：，在这里，我们只看Top N的两个函数，其他函数在用到的时候查看文档就好了。

1）、heapq.nlargest(n, iterable[, key])

从迭代器对象iterable中返回前n个最大的元素列表，其中关键字参数key用于匹配是字典对象的iterable，用于更复杂的数据结构中。

2）、heapq.nsmallest(n, iterable[, key])

从迭代器对象iterable中返回前n个最小的元素列表，其中关键字参数key用于匹配是字典对象的iterable，用于更复杂的数据结构中。

关于第三个参数的应用，我们来看一个例子就明白了。

1 >>> portfolio = [ 2     {
   'name': 'IBM', 'shares': 100, 'price': 91.1}, 3     {
   'name': 'AAPL', 'shares': 50, 'price': 543.22}, 4     {
   'name': 'FB', 'shares': 200, 'price': 21.09}, 5     {
   'name': 'HPQ', 'shares': 35, 'price': 31.75}, 6     {
   'name': 'YHOO', 'shares': 45, 'price': 16.35}, 7     {
   'name': 'ACME', 'shares': 75, 'price': 115.65} 8 ] 9 ... ... ... ... ... ... ... >>> 10 >>> cheap = heapq.nsmallest(3, portfolio, key=lambda s: s['price'])11 >>> print cheap12 [{
   'price': 16.35, 'name': 'YHOO', 'shares': 45}, {
   'price': 21.09, 'name': 'FB', 'shares': 200}, {
   'price': 31.75, 'name': 'HPQ', 'shares': 35}]13 >>> expensive = heapq.nlargest(3, portfolio, key=lambda s: s['price'])14 >>> print expensive15 [{
   'price': 543.22, 'name': 'AAPL', 'shares': 50}, {
   'price': 115.65, 'name': 'ACME', 'shares': 75}, {
   'price': 91.1, 'name': 'IBM', 'shares': 100}]16 >>>

从例子中可以看出，key匹配了portfolio中关键字为‘price’的一行。

到此为止，关于如何应用heapq来求Top N问题，相比通过上面的例子讲解，已经较为熟悉了。现在有几个需要注意的地方：

1）heapq.heapify(iterable)：可以将一个列表转换成heapq

2）在Top N问题中，如果N=1，则直接用max(iterable)/min(iterable)即可。

3）如果N很大，接近集合元素，则为了提高效率，采用sort+切片的方式会更好，如：

求最大的N个元素：sorted(iterable, key=key, reverse=True)[:N]

求最小的N个元素：sorted(iterable, key=key)[:N]

1 >>> nums = [1,8,2,23,7,-4,18,23,42,37,2]2 >>> max(nums)3 424 >>> min(nums)5 -46 >>> print sorted(nums, reverse=True)[:3]7 [42, 37, 23]8 >>> print sorted(nums)[:3]9 [-4, 1, 2]