Python精进系列：heapq 模块

---

---

引言

Python heapq 模块是一个处理堆数据结构的强大工具。堆这种数据结构，以其独特的二叉树特性，在众多算法和数据处理场景中扮演着关键角色。heapq模块默认实现的是最小堆，即父节点的值总是小于或等于其子节点的值。接下来，让我们深入探索heapq的各种用法。

一、heapq基础操作

（一）将列表转换为堆

在实际编程中，我们常常需要将已有的数据结构转换为堆，以便利用堆的特性进行高效处理。heapq模块提供了heapify()函数，它能原地将一个列表转换为堆结构，时间复杂度为O(n)，这意味着对于大规模数据的转换也能高效完成。

import heapq

nums = [3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5]
heapq.heapify(nums)
print(nums)

运行上述代码，输出结果为[1, 1, 2, 3, 3, 9, 4, 6, 5, 5, 5] 。可以看到，原本无序的列表nums已成功转换为一个最小堆，堆顶元素（即列表的第一个元素）为最小值。

（二）向堆中添加元素

向堆中添加元素是常见操作之一。heapq.heappush()函数专门用于此目的，它在保持堆性质的前提下，将新元素添加到堆中。该操作的时间复杂度为O(log n)，其中n是堆中元素的数量。这一特性使得在处理大规模堆时，添加元素的操作依然高效。

import heapq

heap = [1, 3, 5]
heapq.heappush(heap, 2)
print(heap)

运行结果为[1, 2, 5, 3] 。元素2被顺利添加到堆中，且堆的最小堆性质得以维持。

（三）从堆中移除最小元素

当我们需要获取并移除堆中的最小元素时，heapq.heappop()函数便能派上用场。该函数移除并返回堆中的最小元素，同时调整堆结构以保持堆性质，时间复杂度同样为O(log n)。

import heapq

heap = [1, 2, 3, 4, 5]
min_value = heapq.heappop(heap)
print(min_value)
print(heap)

上述代码输出1和[2, 4, 3, 5] 。最小元素1被成功移除并返回，堆结构也相应调整，新的最小元素2位于堆顶。

（四）获取堆中的最小元素

有时，我们仅需查看堆中的最小元素，而不希望对堆结构进行修改。此时，直接访问堆列表的第一个元素即可，因为堆的第一个元素始终是最小元素。

import heapq

heap = [1, 3, 5, 7, 9]
min_value = heap[0]
print(min_value)

输出结果为1 ，通过这种简单方式，我们能快速获取堆中的最小值，且不会改变堆的结构。

二、heapq进阶用法

（一）heapq.heapreplace()的使用

heapq.heapreplace()函数将移除堆中最小元素与添加新元素这两个操作合并为一个原子操作。它先移除堆中的最小元素，然后将指定元素添加到堆中，并返回被移除的最小元素。相较于先调用heappop()再调用heappush()，该函数能在一定程度上提高效率。

import heapq

heap = [1, 3, 5, 7]
replaced_value = heapq.heapreplace(heap, 4)
print(replaced_value)
print(heap)

运行结果为1和[3, 4, 5, 7] 。最小元素1被移除并返回，同时元素4被添加到堆中，堆结构保持最小堆性质。

（二）heapq.merge()合并多个已排序迭代器

在处理多个已排序的数据集时，heapq.merge()函数非常实用。它能合并多个已排序的迭代器（如列表、元组等），并返回一个新的迭代器，该迭代器按升序生成所有输入迭代器中的元素。合并过程借助堆数据结构，时间复杂度为O(N log k)，其中N是所有输入迭代器中元素的总数，k是输入迭代器的数量。

import heapq

list1 = [1, 4, 7]
list2 = [2, 5, 8]
list3 = [3, 6, 9]
merged = list(heapq.merge(list1, list2, list3))
print(merged)

上述代码输出[1, 2, 3, 4, 5, 6, 7, 8, 9] 。通过heapq.merge()函数，三个已排序的列表被高效合并成一个新的已排序列表。

（三）heapq.nlargest()和heapq.nsmallest()获取最值元素

heapq.nlargest()和heapq.nsmallest()函数用于获取堆或可迭代对象中的前n个最大和最小元素。这两个函数在处理大数据集时优势明显，因为它们无需对整个数据集进行排序，而是利用堆的特性快速定位目标元素。

import heapq

nums = [3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5]
largest_three = heapq.nlargest(3, nums)
smallest_two = heapq.nsmallest(2, nums)
print(largest_three)
print(smallest_two)

运行结果为[9, 6, 5]和[1, 1] 。heapq.nlargest(3, nums)返回列表nums中最大的三个元素，heapq.nsmallest(2, nums)返回最小的两个元素。

三、heapq在实际场景中的应用

（一）优先队列的实现

优先队列在许多算法和系统中至关重要，它按照元素的优先级进行排序，优先级高的元素先出队。借助heapq模块，我们能轻松实现优先队列。将元素及其优先级作为元组放入堆中，即可满足优先队列的需求。

import heapq


class PriorityQueue:
    def __init__(self):
        self.heap = []
        self.count = 0

    def push(self, item, priority):
        heapq.heappush(self.heap, (-priority, self.count, item))
        self.count += 1

    def pop(self):
        _, _, item = heapq.heappop(self.heap)
        return item


pq = PriorityQueue()
pq.push('task1', 3)
pq.push('task2', 1)
pq.push('task3', 2)
print(pq.pop())
print(pq.pop())
print(pq.pop())

上述代码定义了一个PriorityQueue类，使用heapq实现优先队列。push方法将任务及其优先级添加到堆中，pop方法移除并返回优先级最高的任务。输出结果为task2、task3、task1，符合优先队列按照优先级从高到低出队的要求。

（二）数据流中位数的计算

在处理数据流时，实时计算中位数是一个常见需求。利用heapq模块，通过维护一个最大堆和一个最小堆，能高效地实现这一功能。合理分配数据流中的元素到两个堆中，即可快速计算中位数。

import heapq


class MedianFinder:
    def __init__(self):
        self.small_heap = []  # 最大堆
        self.large_heap = []  # 最小堆

    def addNum(self, num):
        if not self.small_heap or num <= -self.small_heap[0]:
            heapq.heappush(self.small_heap, -num)
        else:
            heapq.heappush(self.large_heap, num)
        if len(self.small_heap) > len(self.large_heap) + 1:
            heapq.heappush(self.large_heap, -heapq.heappop(self.small_heap))
        elif len(self.large_heap) > len(self.small_heap):
            heapq.heappush(self.small_heap, -heapq.heappop(self.large_heap))

    def findMedian(self):
        if len(self.small_heap) == len(self.large_heap):
            return (-self.small_heap[0] + self.large_heap[0]) / 2
        else:
            return -self.small_heap[0]


mf = MedianFinder()
mf.addNum(1)
mf.addNum(2)
print(mf.findMedian())
mf.addNum(3)
print(mf.findMedian())

这段代码定义了MedianFinder类，通过两个堆（small_heap为最大堆，large_heap为最小堆）维护数据流。addNum方法将新元素添加到合适的堆中，并保证两个堆的大小平衡。findMedian方法根据两个堆的大小关系计算并返回中位数。

（三）Dijkstra算法的实现

Dijkstra算法是寻找加权图中最短路径的经典算法，在实现过程中，需要高效获取当前距离源点最近的节点。heapq模块提供的堆数据结构恰好满足这一需求，能显著提升算法的执行效率。

import heapq


def dijkstra(graph, start):
    distances = {node: float('inf') for node in graph}
    distances[start] = 0
    pq = [(0, start)]
    while pq:
        dist, current = heapq.heappop(pq)
        if dist > distances[current]:
            continue
        for neighbor, weight in graph[current].items():
            new_dist = dist + weight
            if new_dist < distances[neighbor]:
                distances[neighbor] = new_dist
                heapq.heappush(pq, (new_dist, neighbor))
    return distances


graph = {
    'A': {'B': 1, 'C': 4},
    'B': {'A': 1, 'C': 2, 'D': 5},
    'C': {'A': 4, 'B': 2, 'D': 1},
    'D': {'B': 5, 'C': 1}
}
start_node = 'A'
print(dijkstra(graph, start_node))

上述代码实现了Dijkstra算法，利用heapq维护一个优先队列，每次从队列中取出距离源点最近的节点进行扩展。通过这种方式，有效计算出从源点到图中各个节点的最短距离。

四、总结

heapq模块作为Python标准库的重要组成部分，为开发者提供了一套高效处理堆数据结构的工具。从基础的堆创建、元素操作，到进阶的合并、获取最值等功能，再到在优先队列、中位数计算、最短路径算法等实际场景中的广泛应用，heapq展现出了强大的实用性和高效性。在日常编程中，尤其是处理大规模数据时，合理运用heapq模块能够显著提升程序的性能，优化代码结构。