使用 Python 实现优先级队列

本文最后更新于:2021年11月14日 凌晨

有时候会不可避免的用到优先级队列,根据优先级来决定调度,这一点在事件监听器中非常常见。

如何用 Python 实现优先级队列?

Python 有个标准库模块 heapq,这个模块提供了堆队列的相关操作。

简单来说,就是利用了 heapq 的堆队列,构造了一个优先级队列。并利用迭代来逐个获取元素,从而达到 priority 越小越优先的效果。

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
from typing import NoReturn, Any, NewType, Iterator
import heapq

class PriorityQueue(Iterator):
T = NewType('T', Any) # Type[T]

def __init__(self):
self.queue = []
self.index = 0

def join(self, item: T, priority: int) -> NoReturn:
# 使用 heapq.heappush 将 item 插入堆队列
# 如果把 priority 改为 -priority, 可实现 `priority` 越大越优先
heapq.heappush(self.queue, (priority, self.index, item))
# index 索引可以将有相同优先级的 item 以适当的顺序排列, 避免顺序错乱
self.index += 1

def __next__(self) -> T:
# 构造迭代器
# 可以使用 for item in priority_queue 迭代获取最小元素
if self.queue:
# 使用 heapq.heappop 将堆队列中的最小元素弹出
# heapq.heappop 的结果是个元组, 根据 (priority, index, item) 获取 item
# 即获取 tuple[-1]
return heapq.heappop(self.queue)[-1]
else:
raise StopIteration

@property
def first(self) -> T:
# 如果不想用迭代器, 也可以逐个获取 first
# 使用 priority_queue.first 获取
return heapq.heappop(self.queue)[-1]

本博客所有文章除特别声明外,均采用 CC BY-SA 4.0 协议 ,转载请注明出处!