1.5 实现一个优先级队列

1.5 实现一个优先级队列

问题描述

怎样实现一个按优先级排序的队列?并且每次pop操作总是返回优先级最高的那个元素?

解决方案

下面利用heapq模块实现了一个简单的优先级队列类:

import heapq


class PriorityQueue:
    """
    简单的优先级队列类
    """
    def __init__(self):
        self._queue = []
        self._index = 0

    def push(self, item, priority: int):
        """
        在队列中加入新的值,并保持堆的不变性(即保持最小堆)

        :param item: 要加入的值
        :param priority: 优先级
        :return: None
        """
        heapq.heappush(self._queue, (-priority, self._index, item))
        self._index += 1

    def pop(self):
        """
        弹出self._queue[0][-1],即优先级最高的元素的值

        :return: item
        """
        return heapq.heappop(self._queue)[-1]

下面是使用方式:

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

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


q = PriorityQueue()
q.push(Item("happy"), 1)
q.push(Item("good"), 5)
q.push(Item("beautiful"), 3)
q.push(Item("nice"), 1)

for i in range(0, 4):
    print(q.pop())
"""
输出结果:
Item("good")
Item("beautiful")
Item("happy")
Item("nice")
"""

讨论

这里的底层实现是堆排序。函数heapq.heapqpush()heapq.heappop()分别在队列_queue上插入和删除第一个元素,并且因为_queue是最小堆,所以-priority最小的元组总是排在第一个。这样就实现了按优先级返回元素。

在上面的代码中,队列中的每个元素是一个元组(-priority, index, item)。元组的大小比较是从第一个元素开始比较,比较结果即为元组比较结果,如相等则依次向后比较。index变量的作用是保证同优先级的元素按照它们的插入顺序排序和比较。
我们设定Item实例是不支持排序的:

a = Item("happy")
b = Item("good")
a < b
"""
运行结果:
TypeError: "<" not supported between instances of "Item" and "Item"
"""

如果使用元组(-priority, item),那么当优先级相同时,就会出现如上的运行错误。
通过引入另外的index变量组成三元组(-priority, index, item),就能很好的避免此错误,因为不可能有相同的index值。

关于heapq模块的更多详细说明可以参考官方文档