heapq python3.8 官网 (opens new window)

概念

提示

这个模块提供了堆队列算法的实现，也称为优先队列算法

堆是一个二叉树，它的每个父节点的值都只会小于或等于所有孩子节点（的值）。 它使用了数组来实现：从零开始计数，对于所有的 k ，都有 heap[k] <= heap[2k+1] 和 heap[k] <= heap[2k+2]。 为了便于比较，不存在的元素被认为是无限大。

堆最有趣的特性在于最小的元素总是在根结点：heap[0]

这个 API 与教材的堆算法实现有所不同，具体区别有两方面：

（a）我们使用了从零开始的索引。这使得节点和其孩子节点索引之间的关系不太直观但更加适合，因为 Python 使用从零开始的索引。

（b）我们的 pop 方法返回最小的项而不是最大的项（这在教材中称为“最小堆”；而“最大堆”在教材中更为常见，因为它更适用于原地排序）

基于这两方面，把堆看作原生的 Python list 也没什么奇怪的： heap[0] 表示最小的元素，同时 heap.sort() 维护了堆的不变性！

要创建一个堆，可以使用 list 来初始化为 [] ，或者你可以通过一个函数 heapify() ，来把一个 list 转换成堆。

heapq 常用函数

# heapq.heappush(heap, item) ，将 item 的值加入 heap 中，保持堆的不变性
# heapq.heappop(heap) ，弹出并返回 heap 的最小的元素，保持堆的不变性。如果堆为空，抛出 IndexError 。使用 heap[0] ，可以只访问最小的元素而不弹出它
# heapq.heappushpop(heap, item) ，将 item 放入堆中，然后弹出并返回 heap 的最小元素。该组合操作比先调用 heappush() 再调用 heappop() 运行起来更有效率
# heapq.heapify(x) ，将list x 转换成堆，原地，线性时间内
# heapq.heapreplace(heap, item) ，弹出并返回 heap 中最小的一项，同时推入新的 item。 堆的大小不变。 如果堆为空则引发 IndexError，
# 这个单步骤操作比 heappop() 加 heappush() 更高效，并且在使用固定大小的堆时更为适宜。 pop/push 组合总是会从堆中返回一个元素并将其替换为 item
# heapq.merge(*iterables, key=None, reverse=False) ，将多个已排序的输入合并为一个已排序的输出（例如，合并来自多个日志文件的带时间戳的条目）。 返回已排序值的 iterator
# heapq.nlargest(n, iterable, key=None) ，从 iterable 所定义的数据集中返回前 n 个最大元素组成的列表。 如果提供了 key 则其应指定一个单参数的函数，
# 用于从 iterable 的每个元素中提取比较键 (例如 key=str.lower)。 等价于: sorted(iterable, key=key, reverse=True)[:n]
# heapq.nsmallest(n, iterable, key=None) ，从 iterable 所定义的数据集中返回前 n 个最小元素组成的列表。 如果提供了 key 则其应指定一个单参数的函数，
# 用于从 iterable 的每个元素中提取比较键 (例如 key=str.lower)。 等价于: sorted(iterable, key=key)[:n]

# 后两个函数在 n 值较小时性能最好。 对于更大的值，使用 sorted() 函数会更有效率。 此外，当 n==1 时，使用内置的 min() 和 max() 函数会更有效率。
# 如果需要重复使用这些函数，请考虑将可迭代对象转为真正的堆。

import heapq
# 堆排序 可以通过将所有值推入堆中然后每次弹出一个最小值项来实现
import itertools


def heapsort(iterable):
    h = []
    for value in iterable:
        heapq.heappush(h, value)
    return [heapq.heappop(h) for _ in range(len(h))]


print(heapsort([1, 3, 5, 7, 9, 2, 4, 6, 8, 0]))  # [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
# 这类似于 sorted(iterable)，但与 sorted() 不同的是这个实现是不稳定的

# ------------
# 堆元素可以为元组。 这适用于将比较值（例如任务优先级）与跟踪的主记录进行赋值的场合:
h1 = []
heapq.heappush(h1, (5, 'write code'))
heapq.heappush(h1, (7, 'release product'))
heapq.heappush(h1, (1, 'write spec'))
heapq.heappush(h1, (3, 'create tests'))
print(h1)  # [(1, 'write spec'), (3, 'create tests'), (5, 'write code'), (7, 'release product')]
print(heapq.heappop(h1))  # (1, 'write spec')
print(h1)  # [(3, 'create tests'), (7, 'release product'), (5, 'write code')]

# ------------
# 优先队列实现说明
# 优先队列 是堆的常用场合，并且它的实现包含了多个挑战：
# 排序稳定性：你该如何令相同优先级的两个任务按它们最初被加入时的顺序返回？
# 如果优先级相同且任务没有默认比较顺序，则 (priority, task) 对的元组比较将会中断。
# 如果任务优先级发生改变，你该如何将其移至堆中的新位置？
# 或者如果一个挂起的任务需要被删除，你该如何找到它并将其移出队列？

# 针对前两项挑战的一种解决方案是将条目保存为包含优先级、条目计数和任务对象 3 个元素的列表。 条目计数可用来打破平局，
# 这样具有相同优先级的任务将按它们的添加顺序返回。 并且由于没有哪两个条目计数是相同的，元组比较将永远不会直接比较两个任务


# 移除条目或改变其优先级的操作实现起来更为困难，因为它会破坏堆结构不变量。
# 因此，一种可能的解决方案是将条目标记为已移除，再添加一个改变了优先级的新条目：
pq = []  # list of entries arranged in a heap
entry_finder = {}  # mapping of tasks to entries
REMOVED = '<removed-task>'  # placeholder for a removed task
counter = itertools.count()  # unique sequence count


def add_task(task, priority=0):
    'Add a new task or update the priority of an existing task'
    if task in entry_finder:
        remove_task(task)
    count = next(counter)
    entry = [priority, count, task]
    entry_finder[task] = entry
    heapq.heappush(pq, entry)


def remove_task(task):
    'Mark an existing task as REMOVED.  Raise KeyError if not found.'
    entry = entry_finder.pop(task)
    entry[-1] = REMOVED


def pop_task():
    'Remove and return the lowest priority task. Raise KeyError if empty.'
    while pq:
        priority, count, task = heapq.heappop(pq)
        if task is not REMOVED:
            del entry_finder[task]
            return task
    raise KeyError('pop from an empty priority queue')

# ------------
# 理论
# 堆是通过数组来实现的，其中的元素从 0 开始计数，对于所有的 k 都有 a[k] <= a[2*k+1] 且 a[k] <= a[2*k+2]。
# 为了便于比较，不存在的元素被视为无穷大。 堆最有趣的特性在于 a[0] 总是其中最小的元素。