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
模块的更多详细说明可以参考官方文档