[Python知识库]【Python 每日一技】实现一个优先级队列

1 问题

你希望实现一个队列，该队列通过给定的优先级进行元素排序，且每调用一次 pop() 操作都会删除并返回最高优先级的元素。

2. 解决方案

下面是使用 headq 模块的相关函数实现的一个优先级队列及对应测试案例：

import heapq


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

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


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]


def 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())  # Item('bar')
    print(q.pop())  # Item('spam')
    print(q.pop())  # Item('foo')
    print(q.pop())  # Item('grok')


if __name__ == '__main__':
    main()

可以看到，第一个 pop() 操作返回了优先级最高的元素，同时针对具有相同优先级的 foo 和 grok ，在不断调用 pop() 方法时，二者弹出的顺序和插入的顺序相同，即二者优先级虽然相同，但是先插入的元素先弹出，后插入的元素后弹出。

3. 讨论

上述实现的核心是使用了 heapq 模块，具体地，函数 headq.heappush() 和 heapq.heappop() 在实现向列表 self._queue 中插入和删除元素的同时，可以保证列表第一个元素始终是最低优先级的。需要注意的是，这两个函数的最坏时间复杂度都是 $O(\text{log}N)$ ，这意味着即使 $N$ 很大，这两个函数的效率依旧可以很高。

很重要的是，在上述实现中，列表 self._queue 中保存的每个元素都是一个具有三个元素的元组 (-priority, self._index, item) ，其中：

• 表示元素优先级 priority 此处取了相反数，这是为了能让队列中的元素在弹出时可以按照优先级从大到小弹出，否则，在默认情况下元素按照 priority 大小从小到大弹出；
• 实例属性 self._index 用于处理当两个元素的优先级相同时，在调用函数 heapq.heappush() 和 heapq.heappop() 时可以正确地进行元素大小比较。

对于上述第二点，这里做详细说明。首先，两个 Item 类的实例本身是无法直接进行比较的，例如：

>>> class Item:
...    def __init__(self, name):
...        self.name = name
...    def __repr__(self):
...        return 'Item({!r})'.format(self.name)
    
>>> a = Item('foo')
>>> b = Item('bar')
>>> a < b
Traceback (most recent call last):
  File "<input>", line 1, in <module>
TypeError: '<' not supported between instances of 'Item' and 'Item'

如果将插入列表 self._queue 的元素指定为 (-priority, item) 形式的元组，那么只要两个元素的 priority 值不同，那么二者就可以比较了，但是如果两个元素的 priority 值相同，那么比较还是会失败，例如：

>>> a = (-1, Item('foo'))
>>> b = (-5, Item('bar'))
>>> a > b
True
>>> c = (-1, Item('grok'))
>>> a < c
Traceback (most recent call last):
  File "<input>", line 1, in <module>
TypeError: '<' not supported between instances of 'Item' and 'Item'

因此，通过额外引入 self._index 实例属性，使得元素最终为三元组 (-priority, self._index, item) 就可以解决上述问题，因为虽然元素之间的 priority 值可能相同，但是 self._index 对于任意两个元组都不相同，而 Python 在可以通过元组的前若干个元素判断两个元组的大小关系的情况下，就不会尝试继续比较元组后面的元素。例如：

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