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 更详细的内容，可查看官方文档。