Python 實現一個優先級隊列

實現一個優先級隊列

---

問題

---

如何實現一個按優先級排序的隊列，而且每次執行 pop 返回的是優先級最高的元素？

解決方案

---

這裏引用 Python 提供的 heapq 模塊。

import heapq


class PriorityQueue(object):
    '''實現優先級隊列
    每次執行 pop 操做返回優先級最高的元素
    '''

    def __init__(self):
        # 建立初始隊列，以及索引（確保優先級重複的狀況，可以用前後順序區分）
        self._queue = []
        self._index = 0

    def push(self, item, priority):
        # 這裏插入的第二個參數包括 優先級，索引，
        heapq.heappush(self._queue, (-priority, self._index, item))
        self._index += 1

    def pop(self):
        return heapq.heappop(self._queue)[-1]

方法的使用：

class Item(object):
    def __init__(self, name):
        self.name = name

    def __repr__(self):
        return 'Item({!r})'.format(self.name)


q = PriorityQueue()
q.push(Item('foo'), 1)
q.push(Item('bar'), 5)
q.push(Item('spam'), 4)
q.push(Item('grok'), 1)

print(q.pop())
# Item('bar')
print(q.pop())
# Item('spam')
print(q.pop())
# Item('foo')
print(q.pop())
# Item('grok')

這裏能夠觀察到，第一個 pop() 返回的是優先級最高的元素。另外，相同優先級的元素*例子中的 foo 和 grok），pop 操做執行返回的結果是按照插入隊列的順序。

代碼分析

---

這裏主要說起 heapq 模塊的使用。heapq.heappush() 和 heapq.heappop() 分別在隊列 _queue 上插入和刪除第一個元素，而且隊列 _queue 保持第一個元素擁有最高優先級。heappop() 函數老是返回"最小的"元素，這就是保證隊列 pop 操做返回正確元素的關鍵。

上面的代碼中，隊列包含 (-priority, index, item) 的元組。優先級爲負數的目的是使得元素按照優先級從高到低排序。

index 的做用是保證同等優先級元素的前後排序。這裏保存遞增的 index 變量是爲了確保按照他們插入的順序排序。舉例說明 index 變量的做用。

假設 Item 實例不支持排序：

>>> a = Item('foo')
>>> b = Item('bar')
>>> a < b
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
TypeError: unorderable types: Item() < Item()

若是使用元組 (priority, item)，只要兩個元素的優先級不一樣就能比較。可是若是兩個元素優先級同樣，一樣會出錯:

>>> a = (1, Item('foo'))
>>> b = (5, Item('bar'))
>>> a < b
True
>>> c = (1, Item('grok'))
>>> a < c
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
TypeError: unorderable types: Item() < Item()

這裏引入 index 變量就可以解決這個問題，由於 index 是遞增的，不可能相同。而 Python 元組比較的操做是前面的比較肯定結果就不會對後面的進行比較，因此這裏組成的元組爲 (priority, index, item)：

>>> a = (1, 0, Item('foo'))
>>> b = (5, 1, Item('bar'))
>>> c = (1, 2, Item('grok'))
>>> a < b
True
>>> a < c
True

若是須要在多線程中使用同一個隊列，須要適當添加鎖和信號量的機制。

heapq 更詳細的內容，可查看官方文檔。

歡迎關注微信公衆號《書所集錄》