【Python Cookbook】S01E04 实现优先级队列

问题

一个队列，希望能够以给定的优先级进行元素排序，且每次 pop 操作时都能返回优先级最高的哪个元素。

解决方案

利用 heapq 模块实现一个简单的优先级队列：

import heapq

class PriorityQueue:

    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:

    def __init__(self, name):
        self.name = name

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


if __name__ == '__main__':
    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())
    print(q.pop())
    print(q.pop())
    print(q.pop())

讨论

上面的代码片段核心在于 heapq 模块的使用。

• 函数 heapq.heappush() 实现将元素插入到列表 _queue 中。
• 函数 heapq.heappop() 实现将元素从列表 queue 中移除。

具体实现上，队列以元组 (-priority, index, item) 的形式组成，而由于 heapq 模块只支持最小堆，即队列中的最小元素会被优先弹出。所以为了将优先级高的元素先弹出，我们需要将 priority 值取反（通过添加负号 -）来转换为最小堆中的最小值。这样，当 heapq 尝试找出最小值时，它实际上是在找出原始优先级队列中的最大值。

变量 index ，作用则是将具有相同优先级的元素以适当的顺序排列：

# 代码1，包含index
a = (1, 0, Item('foo'))
b = (5, 1, Item('bar'))
c = (1, 2, Item('grok'))
print(a < b)
print(a < c)

# 代码2，不包含index
a = (1, Item('foo'))
b = (5, Item('bar'))
c = (1, Item('grok'))
print(a < b)
print(a < c)

对比上下两份代码，分别包含 index 以及不包含 index，代码2无法对比 $a$ 与 $c$ 的次序，而代码1则可以按照 priority 和 index 次序进行排序。