從一個集合中查詢最大最小的N個元素——Python heapq 堆資料結構

weixin_33858249發表於2015-09-14
Python資料結構

Top N問題在搜尋引擎、推薦系統領域應用很廣， 如果用我們較為常見的語言，如C、C++、Java等，程式碼量至少也得五行，但是用Python的話，只用一個函式就能搞定，只需引入heapq(堆佇列)這個資料結構即可。今天偶然看到這個庫，特意記下之。

先看一個例子：

1 >>> import heapq
2 >>> 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之類的操作，詳細請看官方文件：https://docs.python.org/2/library/heapq.html，在這裡，我們只看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 cheap
12 [{'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 expensive
15 [{'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 42
4 >>> min(nums)
5 -4
6 >>> print sorted(nums, reverse=True)[:3]
7 [42, 37, 23]
8 >>> print sorted(nums)[:3]
9 [-4, 1, 2]

 

相關文章