Python 算法基础篇之堆和优先队列

1. 堆的本质：完全二叉树的魔法

1.1 什么是堆？

堆（Heap） ：一种特殊的完全二叉树 ，满足堆属性------任意节点的值总是不大于（或不小于）其子节点的值。

两种类型：

• 最小堆（Min Heap）：父节点 ≤ 子节点，根节点最小

• 最大堆（Max Heap）：父节点 ≥ 子节点，根节点最大

  最小堆示例： 最大堆示例：

        1                      9
       / \                    / \
      3   2                  7   8
     / \ / \                / \ / \
    6  5 4  7              5  6 3  4

  堆属性：每个父节点 ≤ 子节点 每个父节点 ≥ 子节点

1.2 为什么用完全二叉树？

完全二叉树的数组表示：

树结构：                数组索引：

        0               index:    0   1   2   3   4   5   6
       / \              value:    [1,  3,  2,  6,  5,  4,  7]
      1   2
     / \ / \
    3  4 5  6

父子节点关系（从零开始索引）：
- 父节点 i 的左子节点：2i + 1
- 父节点 i 的右子节点：2i + 2
- 子节点 i 的父节点：(i - 1) // 2

优势：

• 无需指针，数组直接存储
• 插入删除 O(log n)
• 查询最值 O(1)

---

2. 从零实现最小堆

2.1 核心操作

class MinHeap:
    """最小堆的完整实现"""
    
    def __init__(self):
        self.heap = []
    
    def parent(self, i: int) -> int:
        return (i - 1) // 2
    
    def left_child(self, i: int) -> int:
        return 2 * i + 1
    
    def right_child(self, i: int) -> int:
        return 2 * i + 2
    
    def insert(self, k) -> None:
        """
        插入元素：先放末尾，再上浮调整
        时间复杂度：O(log n)
        """
        self.heap.append(k)
        i = len(self.heap) - 1
        
        # 上浮：如果比父节点小，交换上移
        while i > 0 and self.heap[i] < self.heap[self.parent(i)]:
            self.heap[i], self.heap[self.parent(i)] = \
                self.heap[self.parent(i)], self.heap[i]
            i = self.parent(i)
    
    def heapify(self, i: int) -> None:
        """
        下沉调整：以 i 为根的子树不满足堆属性时，向下调整
        时间复杂度：O(log n)
        """
        smallest = i
        left = self.left_child(i)
        right = self.right_child(i)
        n = len(self.heap)
        
        # 找三者中最小
        if left < n and self.heap[left] < self.heap[smallest]:
            smallest = left
        if right < n and self.heap[right] < self.heap[smallest]:
            smallest = right
        
        # 如果最小不是根，交换并递归下沉
        if smallest != i:
            self.heap[i], self.heap[smallest] = \
                self.heap[smallest], self.heap[i]
            self.heapify(smallest)
    
    def extract_min(self):
        """
        提取最小值：移除根节点，末尾元素补位，再下沉调整
        时间复杂度：O(log n)
        """
        if not self.heap:
            return None
        
        root = self.heap[0]
        
        # 用末尾元素替换根
        self.heap[0] = self.heap[-1]
        self.heap.pop()
        
        # 下沉调整
        if self.heap:
            self.heapify(0)
        
        return root
    
    def peek(self):
        """查看最小值，O(1)"""
        return self.heap[0] if self.heap else None
    
    def size(self) -> int:
        return len(self.heap)
    
    def is_empty(self) -> bool:
        return len(self.heap) == 0


# 测试
heap = MinHeap()
for x in [3, 1, 5, 2, 4]:
    heap.insert(x)
    print(f"插入 {x} 后: {heap.heap}")

print(f"\n最小值: {heap.peek()}")  # 1

while not heap.is_empty():
    print(f"弹出: {heap.extract_min()}, 剩余: {heap.heap}")

输出：

插入 3 后: [3]
插入 1 后: [1, 3]
插入 5 后: [1, 3, 5]
插入 2 后: [1, 2, 5, 3]
插入 4 后: [1, 2, 5, 3, 4]

最小值: 1
弹出: 1, 剩余: [2, 3, 5, 4]
弹出: 2, 剩余: [3, 4, 5]
弹出: 3, 剩余: [4, 5]
弹出: 4, 剩余: [5]
弹出: 5, 剩余: []

2.2 插入过程可视化

插入 2 到堆 [1, 3, 5]：

步骤1：放末尾
        1              1
       / \     →      / \
      3   5          3   5
                      /
                     2

步骤2：上浮（2 < 3，交换）
        1              1
       / \     →      / \
      2   5          2   5
     /              /
    3              3

步骤3：2 > 1，停止。堆属性恢复！

2.3 提取最小值可视化

提取堆 [1, 2, 5, 3, 4] 的最小值：

步骤1：记录根 1，末尾 4 补位
        1              4
       / \     →      / \
      2   5          2   5
     / \            / \
    3   4          3   (移除)

步骤2：下沉（4 > 2，交换）
        4              2
       / \     →      / \
      2   5          4   5
     /              /
    3              3

步骤3：4 > 3，交换
        2              2
       / \     →      / \
      4   5          3   5
     /              /
    3              4

堆恢复：[2, 3, 5, 4]

---

3. Python heapq 模块实战

3.1 heapq 核心函数

import heapq

# heapq 操作的是列表，原地修改

data = [3, 1, 5, 2, 4]

# 建堆：O(n)
heapq.heapify(data)
print(data)  # [1, 2, 5, 3, 4]（满足堆属性，但不一定有序）

# 插入：O(log n)
heapq.heappush(data, 0)
print(data)  # [0, 2, 1, 3, 4, 5]

# 弹出最小值：O(log n)
min_val = heapq.heappop(data)
print(min_val, data)  # 0 [1, 2, 5, 3, 4]

# 弹出最小并插入新值（比分开操作快）
heapq.heapreplace(data, 6)  # 先 pop 再 push
print(data)  # [2, 3, 5, 6, 4]

# 先 push 再 pop（保证堆大小不变）
heapq.heappushpop(data, 1)
print(data)

3.2 最大堆的实现

Python 的 heapq 只支持最小堆，实现最大堆有两种方法：

import heapq

# 方法1：存储负数
class MaxHeap:
    def __init__(self):
        self.heap = []
    
    def push(self, val):
        heapq.heappush(self.heap, -val)
    
    def pop(self):
        return -heapq.heappop(self.heap)
    
    def peek(self):
        return -self.heap[0] if self.heap else None


# 方法2：自定义比较（更通用）
class Item:
    def __init__(self, priority, task):
        self.priority = priority
        self.task = task
    
    def __lt__(self, other):
        # 反转比较，实现最大堆
        return self.priority > other.priority
    
    def __repr__(self):
        return f"Item({self.priority}, {self.task})"


# 测试
max_heap = MaxHeap()
max_heap.push(3)
max_heap.push(1)
max_heap.push(5)
print(max_heap.pop())  # 5

---

4. 经典算法问题

4.1 Top K 问题

问题：从 n 个数中找出最大的 k 个。

import heapq

def top_k_largest(nums: list[int], k: int) -> list[int]:
    """
    找最大的 k 个数
    
    思路：维护大小为 k 的最小堆
          堆顶是第 k 大的，比它大的替换进来
    时间：O(n log k)，空间：O(k)
    """
    if k <= 0:
        return []
    if k >= len(nums):
        return sorted(nums, reverse=True)
    
    # 前 k 个建堆
    heap = nums[:k]
    heapq.heapify(heap)
    
    # 遍历剩余元素
    for num in nums[k:]:
        if num > heap[0]:           # 比堆顶（第k大）大
            heapq.heapreplace(heap, num)
    
    # 返回排序后的结果
    return sorted(heap, reverse=True)


# 测试
nums = [3, 2, 1, 5, 6, 4]
print(top_k_largest(nums, 2))  # [6, 5]

nums = [3, 2, 3, 1, 2, 4, 5, 5, 6]
print(top_k_largest(nums, 4))  # [6, 5, 5, 4]

4.2 合并 K 个有序链表

import heapq

class ListNode:
    def __init__(self, val=0, next=None):
        self.val = val
        self.next = next
    
    def __lt__(self, other):
        # 用于堆比较
        return self.val < other.val

def merge_k_lists(lists: list[ListNode]) -> ListNode:
    """
    合并 k 个升序链表
    
    思路：最小堆存各链表头节点，每次取最小，接上新节点
    时间：O(N log k)，N 为总节点数，k 为链表数
    """
    dummy = ListNode(0)
    current = dummy
    
    # 初始化堆（过滤空链表）
    heap = [node for node in lists if node]
    heapq.heapify(heap)
    
    while heap:
        # 取出最小节点
        node = heapq.heappop(heap)
        current.next = node
        current = current.next
        
        # 该链表还有节点，继续入堆
        if node.next:
            heapq.heappush(heap, node.next)
    
    return dummy.next

4.3 数据流的中位数

import heapq

class MedianFinder:
    """
    动态维护数据流的中位数
    
    思路：
    - 大顶堆（左半部分）：存较小的一半，堆顶是最大值
    - 小顶堆（右半部分）：存较大的一半，堆顶是最小值
    - 保持两个堆大小平衡（差不超过1）
    """
    
    def __init__(self):
        # 大顶堆：用负数实现
        self.small = []   # 较小的一半（实际存负数）
        # 小顶堆
        self.large = []   # 较大的一半
    
    def add_num(self, num: int) -> None:
        # 先放入大顶堆
        heapq.heappush(self.small, -num)
        
        # 大顶堆最大值 移到 小顶堆
        val = -heapq.heappop(self.small)
        heapq.heappush(self.large, val)
        
        # 平衡大小：小顶堆最多比大顶堆多1个
        if len(self.large) > len(self.small) + 1:
            val = heapq.heappop(self.large)
            heapq.heappush(self.small, -val)
    
    def find_median(self) -> float:
        if len(self.small) > len(self.large):
            return -self.small[0]
        elif len(self.large) > len(self.small):
            return self.large[0]
        else:
            return (-self.small[0] + self.large[0]) / 2


# 测试
mf = MedianFinder()
for num in [1, 2, 3, 4, 5]:
    mf.add_num(num)
    print(f"添加 {num}，中位数: {mf.find_median()}")

输出：

添加 1，中位数: 1
添加 2，中位数: 1.5
添加 3，中位数: 2
添加 4，中位数: 2.5
添加 5，中位数: 3

4.4 Dijkstra 最短路径

import heapq

def dijkstra(graph: dict, start: str) -> dict:
    """
    Dijkstra 单源最短路径
    
    参数：graph = {节点: {邻居: 权重}}
    返回：{节点: 最短距离}
    """
    # 初始化距离
    dist = {node: float('inf') for node in graph}
    dist[start] = 0
    
    # 优先队列：(距离, 节点)
    pq = [(0, start)]
    
    while pq:
        d, node = heapq.heappop(pq)
        
        # 已找到更短路径，跳过
        if d > dist[node]:
            continue
        
        # 松弛相邻边
        for neighbor, weight in graph[node].items():
            new_dist = d + weight
            if new_dist < dist[neighbor]:
                dist[neighbor] = new_dist
                heapq.heappush(pq, (new_dist, neighbor))
    
    return dist


# 测试
graph = {
    'A': {'B': 4, 'C': 2},
    'B': {'A': 4, 'C': 1, 'D': 5},
    'C': {'A': 2, 'B': 1, 'D': 8, 'E': 10},
    'D': {'B': 5, 'C': 8, 'E': 2, 'F': 6},
    'E': {'C': 10, 'D': 2, 'F': 3},
    'F': {'D': 6, 'E': 3}
}

print(dijkstra(graph, 'A'))
# {'A': 0, 'B': 3, 'C': 2, 'D': 8, 'E': 10, 'F': 13}

---

5. 总结

5.1 堆操作复杂度

操作	时间复杂度	说明
建堆	O(n)	heapify
插入	O(log n)	上浮调整
删除最小	O(log n)	下沉调整
查看最小	O(1)	直接取根
堆排序	O(n log n)	n 次提取

5.2 应用场景速查

场景	解法	关键思路
Top K	大小为 k 的堆	最小堆找最大，最大堆找最小
中位数	双堆	大顶堆+小顶堆，保持平衡
合并有序	堆存各序列当前元素	每次取最小，补充新元素
最短路径	Dijkstra + 优先队列	按距离优先级扩展
任务调度	优先队列	按优先级/截止时间处理

5.3 核心要点

堆的本质：
├── 完全二叉树的数组表示
├── 堆属性：父节点与子节点的大小关系
└── 核心操作：上浮（插入）、下沉（删除）

Python 实践：
├── heapq 模块：最小堆，原地操作列表
├── 最大堆：存负数或自定义 __lt__
└── 优先队列：heapq + 自定义对象

算法套路：
├── Top K = 维护大小为 K 的堆
├── 中位数 = 双堆平衡
└── 多路归并 = 堆存各路人马当前元素