Python PriorityQueue能否按优先级分组获取队列大小？

好问题！原生的PriorityQueue确实没有内置方法来分别统计不同优先级的元素数量，因为它底层是基于堆结构实现的，元素是按堆的规则排列的，并没有维护各优先级的计数信息。不过我们可以通过几种方式来实现这个需求，下面给你详细讲讲：

方法1：手动维护优先级计数器 #

最简单的方式就是在每次存入或取出元素时，手动更新对应优先级的计数。比如用一个字典来记录各优先级的元素数量：

from queue import PriorityQueue

# 初始化队列和计数器
q = PriorityQueue()
priority_counts = {1: 0, 2: 0}

# 存入元素时更新计数
some_variable = True
my_object = "test_object"
if some_variable:
    q.put((1, my_object))
    priority_counts[1] += 1
else:
    q.put((2, my_object))
    priority_counts[2] += 1

# 取出元素时也要更新计数
if not q.empty():
    prio, obj = q.get()
    priority_counts[prio] -= 1

# 查看各优先级的元素数量
print(f"优先级1的元素数量：{priority_counts[1]}")
print(f"优先级2的元素数量：{priority_counts[2]}")

这种方式的优点是实现简单，不需要额外封装类；缺点是需要手动同步计数和队列操作，容易因为遗漏更新导致计数错误。

方法2：封装自定义的带计数功能的优先级队列 #

更优雅的方式是继承PriorityQueue，封装一个自带优先级计数功能的类，把计数逻辑整合到put和get方法中：

from queue import PriorityQueue, Empty

class CountedPriorityQueue(PriorityQueue):
    def __init__(self):
        super().__init__()
        self.priority_counts = {}  # 存储各优先级的元素数量
    
    def put(self, item, block=True, timeout=None):
        prio, _ = item
        # 初始化该优先级的计数（如果不存在）
        if prio not in self.priority_counts:
            self.priority_counts[prio] = 0
        self.priority_counts[prio] += 1
        super().put(item, block, timeout)
    
    def get(self, block=True, timeout=None):
        try:
            prio, obj = super().get(block, timeout)
            self.priority_counts[prio] -= 1
            # 如果计数为0，删除该优先级的键（可选操作）
            if self.priority_counts[prio] == 0:
                del self.priority_counts[prio]
            return prio, obj
        except Empty:
            raise Empty("队列为空")
    
    def get_priority_count(self, prio):
        """获取指定优先级的元素数量"""
        return self.priority_counts.get(prio, 0)

# 使用示例
q = CountedPriorityQueue()
q.put((1, "obj1"))
q.put((2, "obj2"))
q.put((1, "obj3"))

print(q.get_priority_count(1))  # 输出：2
print(q.get_priority_count(2))  # 输出：1

# 取出一个优先级1的元素
q.get()
print(q.get_priority_count(1))  # 输出：1

这种方式把计数逻辑和队列操作绑定在一起，避免了手动更新的失误，还可以轻松扩展支持更多优先级。

方法3：用两个独立队列模拟双优先级队列 #

如果你的场景只需要支持优先级1和2这两种情况，还可以直接用两个普通的Queue分别存储不同优先级的元素，然后封装一个类来模拟优先级队列的行为：

from queue import Queue, Empty

class DualPriorityQueue:
    def __init__(self):
        self.p1_queue = Queue()  # 存储优先级1的元素
        self.p2_queue = Queue()  # 存储优先级2的元素
    
    def put(self, item, prio):
        if prio == 1:
            self.p1_queue.put(item)
        elif prio == 2:
            self.p2_queue.put(item)
        else:
            raise ValueError("仅支持优先级1和2")
    
    def get(self):
        # 优先取出优先级1的元素
        if not self.p1_queue.empty():
            return (1, self.p1_queue.get())
        elif not self.p2_queue.empty():
            return (2, self.p2_queue.get())
        else:
            raise Empty("队列为空")
    
    def get_priority_count(self, prio):
        if prio == 1:
            return self.p1_queue.qsize()
        elif prio == 2:
            return self.p2_queue.qsize()
        else:
            return 0

# 使用示例
q = DualPriorityQueue()
q.put("obj1", 1)
q.put("obj2", 2)
q.put("obj3", 1)

print(q.get_priority_count(1))  # 输出：2
print(q.get_priority_count(2))  # 输出：1

这种方式的优点是性能更高（不需要堆操作），而且统计数量直接调用每个队列的qsize()，非常直观。

内容的提问来源于stack exchange，提问作者10mjg