question: 现在想实现一个能够以给定的优先级来对元素排序,且每次pop操作时都会返回优先级最高的那个元素。
solution: 利用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]
函数heapq模块的使用:函数heapq.heappush()以及heapq.heappop()分别实现将元素从列表_queue中插入和移除,且保证列表中第一个元素的优先级最低。heappop()方法总是返回“最小”元素。这就是让队列能正确弹出正确元素的关键了。push和pop操作的复杂度都是O(logN),N代表堆中元素的数量,所以就算N的值很大,这些操作的效率也非常高。
在上面的代码中, 队列以元组(-priority, self._index, item)的形式组成。把priority取负值是为了让队列能够按元素的优先级从高到低的顺序排列。这和正常的堆排列顺序相反,一般的堆排序都是从小到大的顺序排列的。
变量index的作用是为了将具有相同优先级的元素以适当的顺序排列,通过维护一个不断递增的索引,元素将以它们入队列时的顺序来排列,所以index在对具有相同优先级的元素间做比较时同样扮演了重要的角色。
实例:
class Item:
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('spam'), 2)
q.push(Item('grok'), 3)
q.push(Item('bob'), 5)
# 把数据取出
print(q.pop())
print(q.pop())
print(q.pop())
print(q.pop())
现在来看看上面的例子,第一次执行pop()操作时返回的元素具有最高的优先级,当两个拥有相同优先级的元素(foo和bob)返回的顺序跟它们插入到队列时的顺序相同。