「数据结构和算法」5-实现一个优先级队列

问题描述

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

解决方案

可使用 heapq 模块实现了一个简单的优先级队列。

代码演示:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
import heapq


class PriorityQueue:
"""利用 heapq 模块实现了一个简单的优先级队列"""

def __init__(self):
self._queue = []
self._index = 0

def push(self, item, priority):
"""
添加队列
:param item: 元素
:param priority: 优先级
"""
# 优先级为负数的目的是使得元素按照优先级从高到低排序
# 这个跟普通的按优先级从低到高排序的堆排序恰巧相反
heapq.heappush(self._queue, (-priority, self._index, item))
# 通过不断增加 index 保证元素按照它们插入的顺序挺排序
self._index += 1

def pop(self):
"""弹出队列"""
return heapq.heappop(self._queue)[-1]


class Item:
"""使用方式"""
def __init__(self, name):
self.name = name

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


def test_priority_queue():
q = PriorityQueue()

# 插入元素
q.push(Item('foo'), 1)
q.push(Item('bar'), 5)
q.push(Item('spam'), 4)
q.push(Item('grok'), 1)

# 按优先级 pop 元素
print(q.pop()) # Item('bar')
print(q.pop()) # Item('spam')

# 如果优先级相同,则按插入顺序返回
print(q.pop()) # Item('foo')
print(q.pop()) # Item('grok')

函数 heapq.heappush()heapq.heappop() 分别在队列 _queue 上插入和删除第一个元素, 并且队列 _queue 保证第一个元素拥有最高优先级

heappop() 函数总是返回「最小的」的元素,这就是保证队列 pop 操作返回正确元素的关键。

另外,由于 push 和 pop 操作时间复杂度为 O(log N),其中 N 是堆的大小,因此就算是 N 很大的时候它们运行速度也依旧很快。

在上面代码中,队列包含了一个 (-priority, index, item) 的元组。

优先级为负数-priority)的目的是使得元素按照优先级从高到低排序。 这个跟普通的按优先级从低到高排序的堆排序恰巧相反。

index 变量的作用是保证同等优先级元素的正确排序。 通过保存一个不断增加的 index 下标变量,可以确保元素按照它们插入的顺序排序。 而且, index 变量也在相同优先级元素比较的时候起到重要作用。

Tips:

如果我们想在多个线程中使用同一个队列,就需要增加适当的锁和信号量机制。

毕小烦 wechat
「请扫一扫上面的二维码,关注老毕的微信公众号」
「您的赞赏是老毕持续创作的动力」