Python中heapq與優先隊列【詳細】

本文始發於我的公衆號：TechFlow, 原創不易，求個關注

今天的文章來介紹Python當中一個蠻有用的庫——heapq。

heapq的全寫是heap queue，是堆隊列的意思。這裏的堆和隊列都是數據結構，在後序的文章當中咱們會詳細介紹，今天只介紹heapq的用法，若是不瞭解heap和queue原理的同窗能夠忽略，咱們並不會深刻太多，會在以後的文章裏詳細闡述。

在介紹用法以前，咱們須要先知道優先隊列的定義。隊列你們應該都不陌生，也是很是基礎簡單的數據結構。咱們能夠想象成隊列裏的全部元素排成一排，新的元素只能從隊尾加入隊列，元素要出隊列只能經過隊首，不能中途從隊列當中退出。而優先隊列呢，是給隊列當中的元素每個都設置了優先級，使得隊伍當中的元素會自動按照優先級排序，優先級高的排在前面。

也就是說Python當中的heapq就是一個維護優先隊列的library，咱們經過調用它能夠輕鬆實現優先隊列的功能。

最大或最小的K個元素

咱們來看一個實際的問題，假設咱們當下有N個雜亂無章的元素，可是咱們只關心其中最大的K個或者是最小的K個元素。咱們想從整個數組當中將這部分抽取出來，應該怎麼辦呢？

這個問題在實際當中很是常見，隨便就能夠舉出例子來。好比用戶輸入了搜索詞，咱們根據用戶的搜索詞找到了大量的內容。咱們想要根據算法篩選出用戶最有可能點擊的文原本，機器學習的模型能夠給每個文本一個預測的分數。以後，咱們就須要選出分數最大的K個結果。這種相似的場景還有不少，利用heapq庫裏的nlargest和nsmallest接口能夠很是方便地作到這點。

咱們一塊兒來看一個例子：

import heapq

nums = [14, 20, 5, 28, 1, 21, 16, 22, 17, 28]
heapq.nlargest(3, nums)
# [28, 28, 22]
heapq.nsmallest(3, nums)
# [1, 5, 14]
複製代碼

heapq的nlargest和nsmallest接受兩個參數，第一個參數是K，也就是返回的元素的數量，第二個參數是傳入的數組，heapq返回的正是傳入的數組當中的前K大或者是前K小。

這裏有一個問題，若是咱們數組當中的元素是一個對象呢？應該怎麼辦？

其實也很簡單，有了解過Python自定義關鍵詞排序的同窗應該知道，和排序同樣，咱們能夠經過匿名函數實現。

匿名函數

咱們都知道，在Python當中經過def能夠定義一個函數。經過def定義的函數都有函數名，因此稱爲有名函數。除了有名函數以外，Python還支持匿名函數。顧名思義，就是沒有函數名的函數。也就是說它其餘方面都和普通函數同樣，只不過沒有名字而已。

初學者可能會納悶，函數沒有名字應該怎麼調用呢？

會有這個疑惑很正常，這是由於習慣了面向過程的編程，對面向對象理解不夠深刻致使的。在許多高級語言當中，一切皆對象，一個類，一個函數，一個int都是對象。既然函數也是對象，那麼函數天然也能夠用來傳遞，不只能夠用來傳遞，還能夠用來返回。這是函數式編程的概念了，咱們這裏很少作深刻。

固然，普通函數也同樣能夠傳遞，起到的效果同樣。只不過在編程當中，有些函數咱們只會使用一次，不必再單獨定義一個函數，使用匿名函數會很是方便。

舉個例子，比方說我有一個這樣的函數：

def operate(x, func):
  return func(x)
複製代碼

這個operate函數它接受兩個參數，第一個參數是變量x，第二個參數是一個函數。它會在函數內部調用func，返回func調用的結果。我如今要作這樣一件事情，我但願根據x這個整數對4取餘的餘數來判斷應該用什麼樣的func。若是對4的餘數爲0，我但願求一次方，若是餘數是2，我但願求平方，以此類推。若是按照正常的方法，咱們須要實現4個方法，而後依次傳遞。

這固然是能夠的，不過很是麻煩，若是使用匿名函數，就能夠大大簡化代碼量：

def get_result(x):
  if x % 4 == 0:
    return operate(x, lambda x: x)
  elif x % 4 == 1:
    return operate(x, lambda x: x ** 2)
  elif x % 4 == 2:
    return operate(x, lambda x: x ** 3)
  else:
    return operate(x, lambda x: x ** 4)
複製代碼

在上面的代碼當中，咱們經過lambda關鍵字定義了匿名函數，避免了定義四種函數用來傳遞的狀況。固然，這個問題還有更簡單的寫法，能夠只用一個函數解決。

咱們來看lambda定義匿名函數的語法，首先是lambda關鍵字，表示咱們當下定義的是一個匿名函數。以後跟的是這個匿名函數的參數，咱們只用到一個變量x，因此只須要寫一個x。若是咱們須要用到多個參數，經過逗號分隔，固然也能夠不用參數。寫完參數以後，咱們用冒號分開，冒號後面寫的是返回的結果。

咱們也能夠把匿名函數賦值給一個變量，以後咱們就能夠和調用普通函數同樣來調用了：

square = lambda x: x ** 2

print(square(3))
print(operate(3, square))
複製代碼

自定義排序

回到以前的內容，若是咱們想要heapq排序的是一個對象。那麼heapq並不知道應該依據對象當中的哪一個參數來做爲排序的衡量標準，因此這個時候，須要咱們本身定義一個獲取關鍵字的函數，傳遞給heapq，這樣才能夠完成排序。

好比說，咱們如今有一批電腦，咱們但願heapq可以根據電腦的價格排序：

laptops = [
    {'name': 'ThinkPad', 'amount': 100, 'price': 91.1},
    {'name': 'Mac', 'amount': 50, 'price': 543.22},
    {'name': 'Surface', 'amount': 200, 'price': 21.09},
    {'name': 'Alienware', 'amount': 35, 'price': 31.75},
    {'name': 'Lenovo', 'amount': 45, 'price': 16.35},
    {'name': 'Huawei', 'amount': 75, 'price': 115.65}
]

cheap = heapq.nsmallest(3, portfolio, key=lambda s: s['price'])
expensive = heapq.nlargest(3, portfolio, key=lambda s: s['price'])
複製代碼

在調用nlargest和nsmallest的時候，咱們額外傳遞了一個參數key，咱們傳入的是一個匿名函數，它返回的結果是這個對象的price，也就是說咱們但願heapq根據對象的price來進行排序。

優先隊列

heapq除了能夠返回最大最小的K個數以外，還實現了優先隊列的接口。咱們能夠直接調用heapq.heapify方法，輸入一個數組，返回的結果是根據這個數組生成的堆（等價於優先隊列）。

固然咱們也能夠從零開始，直接經過調用heapq的push和pop來維護這個堆。接下來，咱們就經過heapq來本身動手實現一個優先隊列，代碼很是的簡單，我想你們應該能夠瞬間學會。

首先是實現優先隊列的部分：

import heapq

class PriorityQueue:
  
  def __init__(self):
    self._queue = []
    self._index =0
    
  def push(self, item, priority):
    # 傳入兩個參數，一個是存放元素的數組，另外一個是要存儲的元素，這裏是一個元組。
    # 因爲heap內部默認有小到大排，因此對priority取負數
    heapq.heappush(self._queue, (-priority, self._index, item))
    self._index += 1
  
  def pop(self):
    return heapq.heappop(self._queue)[-1]
複製代碼

其次咱們來實際看一下運用的狀況：

q = PriorityQueue()

q.push('lenovo', 1)
q.push('Mac', 5)
q.push('ThinkPad', 2)
q.push('Surface', 3)

q.pop()
# Mac
q.pop()
# Surface
複製代碼

到這裏，關於heapq的應用方面就算是介紹完了，可是尚未真正的結束。

咱們須要分析一下heapq當中操做的複雜度，關於堆的部分咱們暫時跳過，咱們先來看nlargest和nsmallest。我在github當中找到了這個庫的源碼，在方法的註釋上，做者寫下了這個方法的複雜度，和排序以後取前K個開銷五五開：

def nlargest(n, iterable, key=None):
    """Find the n largest elements in a dataset. Equivalent to: sorted(iterable, key=key, reverse=True)[:n] """
複製代碼

咱們都知道排序的複雜度的指望是，若是你瞭解堆的話，會知道堆一次插入元素的複雜度是。若是咱們限定堆的長度是K，咱們插入n次以後也只能保留K個元素。每次插入的複雜度是，一共插入n次，因此總體的複雜度是。

若是K小一些，可能開銷會比排序稍小，可是程度有限。那麼有沒有什麼辦法能夠不用排序而且儘量快地篩選出前K大或者是前K小的元素呢？

我這裏先賣個關子，咱們以後的文章當中再來說解。

今天的文章就到這裏，若是以爲有所收穫，請順手點個關注吧，你的舉手之勞對我很重要。

參考資料

Python CookBook Version3

維基百科