從一個集合中查詢最大最小的N個元素——Python heapq 堆資料結構
阿新 • • 發佈:2019-01-25
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] 複製程式碼