2.38. 优先队列

优先队列是一种容器,它无需对整体排序就能始终知道其中哪个元素最小(或最大)。heapq 在 MicroPython 中以一个普通的 list 作为存储,并使用一种称为最小堆的基于树的算法,将最小元素保持在索引 0 处。

2.38.1. 思维模型

最小堆是一种这样排列的列表:每个索引处的元素都小于或等于索引 2*i + 12*i + 2 处的元素。这种排列方式使得三种操作的开销都很低:

  • 查找最小元素——它位于 heap[0]。只需一次索引查找,无需扫描。

  • 添加元素——它会沿树向上冒泡,每一层与一个父节点比较。一千个元素大约需要十次比较;一百万个大约需要二十次。

  • 移除最小元素——同样的过程反向进行,比较次数也差不多。

该列表并未排序。遍历它会以任意顺序得到各元素。只有 heap[0] 保证是最小值。

2.38.2. MicroPython 上的三个函数

MicroPython 的 heapq 恰好提供三个函数:

这就是它的全部接口。CPython 的 heappushpopheapreplacenlargestnsmallestmerge存在;下面的范例展示如何手动构建这些常见操作。

从零开始::

>>> import heapq
>>> heap = []
>>> heapq.heappush(heap, 5)
>>> heapq.heappush(heap, 1)
>>> heapq.heappush(heap, 3)
>>> heap
[1, 5, 3]
>>> heapq.heappop(heap)
1
>>> heapq.heappop(heap)
3
>>> heapq.heappop(heap)
5

从现有数据开始::

>>> samples = [5, 1, 3, 2, 4]
>>> heapq.heapify(samples)
>>> samples
[1, 2, 3, 5, 4]
>>> heapq.heappop(samples)
1

堆化只对每个元素处理一次,明显比逐个推入相同的元素要快。

2.38.3. 流中的 Top-N

优先队列的一个自然用途:在只能读取一次的流中追踪N个最大元素。

诀窍是维护一个大小为N的最小堆。每个新值都被推入;如果堆的大小超过N,就丢弃最小的那个。最小元素始终位于 heap[0],因此只需一次比较即可决定保留还是丢弃每个进入的元素::

def top_n_brightest(readings, n):
    heap = []
    for r in readings:
        if len(heap) < n:
            heapq.heappush(heap, r)
        elif r > heap[0]:
            heapq.heappop(heap)
            heapq.heappush(heap, r)
    return sorted(heap, reverse=True)

堆的大小永远不会超过N,因此无论流有多长,每个进入的值都只花费同样几次比较。最后的 sorted 是唯一一个开销直接取决于N的步骤。

2.38.4. 定时事件

(deadline, payload) 元组构成的优先队列是事件调度器的标准形式。元组按字典序比较,因此截止时间最小的元组会自动浮到顶部::

import heapq
import time

queue = []
heapq.heappush(queue, (time.time() + 5, 'check sensors'))
heapq.heappush(queue, (time.time() + 1, 'blink LED'))
heapq.heappush(queue, (time.time() + 3, 'send heartbeat'))

while queue:
    deadline, task = queue[0]
    now = time.time()
    if now < deadline:
        time.sleep(deadline - now)
    heapq.heappop(queue)
    run(task)

如果两个事件的截止时间相同,元组比较会继续比较第二个元素。对于本身不可比较的负载,可在元组的第二个元素位置添加一个计数器以强制稳定排序::

counter = 0
def schedule(deadline, task):
    global counter
    counter += 1
    heapq.heappush(queue, (deadline, counter, task))

现在最小截止时间优先的行为得以保持,平局时按计数器决出先后,从而让更早调度的任务先运行,而 task 本身也无需可比较。

2.38.5. heapq 没有提供什么

有几种看似应该可行、实则不行的模式:

  • for x in heap 遍历不会给出有序输出——列表的底层顺序是堆的树形布局。要按顺序消费堆,请不断 heappop 直到为空。

  • 无法修改堆中已有元素的优先级。标准的变通做法是将新优先级作为一个全新条目推入,并将旧条目标记为已取消(当它弹出时跳过它)。

  • 堆存储的是对其元素的引用,而非副本。在元素推入后对其进行修改可能会破坏堆的不变性。请使用不可变值(数字、元组)作为比较键。

对于小型嵌入式设备上任何"最小是什么"或"最紧急是什么"的问题,基于普通列表的最小堆都是合适的工具。三个函数、一个列表,剩下的就交给堆来处理。