【数据结构】最大最小堆实现优先队列 python

news2025/2/28 15:18:21

堆的定义

堆（Heap）是一种特殊的完全二叉树结构，通常分为最大堆和最小堆两种类型。

在最大堆中，父节点的值总是大于或等于其子节点的值；

而在最小堆中，父节点的值总是小于或等于其子节点的值。

堆常用于实现优先队列，在许多算法中也有重要应用，比如堆排序、Dijkstra算法等。

堆的基本操作

插入：向堆中添加一个新元素，并调整堆以保持其性质。
删除：移除堆顶元素（最大或最小元素），并重新调整堆。
获取最大/最小元素：直接访问堆顶元素即可获得。

堆的实现

Python 的 heapq 模块提供了对堆的支持，它实现了最小堆。以下是一个简单的例子：

import heapq

# 创建一个空堆
heap = []

# 向堆中插入元素
heapq.heappush(heap, 10)
heapq.heappush(heap, 20)
heapq.heappush(heap, 5)
print(heap)

# 获取堆顶元素（最小元素）
min_element = heap[0]
print("堆顶元素:", min_element)

# 移除堆顶元素
heapq.heappop(heap)
print("移除堆顶元素后的堆:", heap)

# 如果需要使用最大堆，可以通过插入负值来模拟
max_heap = []
heapq.heappush(max_heap, -10)
heapq.heappush(max_heap, -20)
heapq.heappush(max_heap, -5)
print(heap)

max_element = -max_heap[0]  # 记得取负数得到原始的最大值
print("最大堆顶元素:", max_element)

注意：为了实现最大堆，我们需要存储元素的负值

因为 Python 标准库中的heapq 模块只提供最小堆的功能。

【模板】堆

题目描述

给定一个数列，初始为空，请支持下面三种操作：

给定一个整数 $x$ ，请将 $x$ 加入到数列中。
输出数列中最小的数。
删除数列中最小的数（如果有多个数最小，只删除 $1$ 个）。

输入格式

第一行是一个整数，表示操作的次数 $n$ 。
接下来 $n$ 行，每行表示一次操作。每行首先有一个整数 $o p$ 表示操作类型。

若 $o p = 1$ ，则后面有一个整数 $x$ ，表示要将 $x$ 加入数列。
若 $o p = 2$ ，则表示要求输出数列中的最小数。
若 $o p = 3$ ，则表示删除数列中的最小数。如果有多个数最小，只删除 $1$ 个。

输出格式

对于每个操作 $2$ ，输出一行一个整数表示答案。

输入输出样例 #1

输入 #1

输出 #1

2
5

说明/提示

【数据规模与约定】

对于 $30\%$ 的数据，保证 $\leq 15$ 。
对于 $70\%$ 的数据，保证 $\leq 10^4$ 。
对于 $100\%$ 的数据，保证 $\leq n \leq 10^6$ ， $\leq x \lt 2^{31}$ ， $\in \{1, 2, 3\}$ 。

AC_code:

import heapq
hp = []
n = int(input())
for _ in range(n):
    a = list(map(int, input().split()))
    op = a[0]
    if op == 1:
        heapq.heappush(hp, a[1])
    elif op == 2:
        print(hp[0])
    else:
        heapq.heappop(hp)

实战演练

牛客周赛 Round 82 E题和+和

在这里插入图片描述

方法思路

我们需要从两个数组中选择2m个下标，使得前m个下标对应的a数组元素之和加上后m个下标对应的b数组元素之和最小。关键在于高效地找到这些下标的最优组合。

预处理前缀和后缀数组：
- pre_a[k]：表示在前k个元素中选择m个最小的a数组元素之和。
- suf_b[k]：表示从第k个元素到末尾中选择m个最小的b数组元素之和。
使用大根堆维护最小元素：
- 遍历数组时维护一个大根堆，确保堆中始终保存当前最小的m个元素。当堆的大小超过m时，弹出最大的元素，保持堆的大小为m。
遍历所有可能的分割点：
- 对于每个可能的分割点k，计算前k个元素的最小m个a之和，加上从k+1到末尾的最小m个b之和，取最小值。

AC_code:

import heapq

n, m = map(int, input().split())
a = list(map(int, input().split()))
b = list(map(int, input().split()))

pre_a = [float('inf')] * (n + 1)
heap = []  # 最大堆（通过存储负数模拟）
sum_a = 0  

for i in range(1, n + 1):
    heapq.heappush(heap, -a[i - 1])  # 将当前元素加入堆中（存入负数表示最大堆）
    sum_a += a[i - 1]  

    # 如果堆的大小超过 m，则移除堆顶元素（即最大的那个）
    if len(heap) > m:
        removed = -heapq.heappop(heap)
        sum_a -= removed

    # 如果堆中有正好 m 个元素，则记录当前的前缀和
    if len(heap) == m:
        pre_a[i] = sum_a

# 初始化后缀和数组 suf_b
suf_b = [float('inf')] * (n + 2)
heap = []  # 清空堆
sum_b = 0 

# 计算后缀和
for i in range(n, 0, -1):
    heapq.heappush(heap, -b[i - 1])  # 将当前元素加入堆中（存入负数表示最大堆）
    sum_b += b[i - 1]  

    # 如果堆的大小超过 m，则移除堆顶元素（即最大的那个）
    if len(heap) > m:
        removed = -heapq.heappop(heap)
        sum_b -= removed

    # 如果堆中有正好 m 个元素，则记录当前的后缀和
    if len(heap) == m:
        suf_b[i] = sum_b

ans = float('inf')
for k in range(m, n-m + 1):
    ans = min(ans, pre_a[k] + suf_b[k+1])

print(ans)

代码解释

预处理前缀数组pre_a：
- 使用大根堆维护前k个元素中最小的m个元素之和。每次添加元素后，若堆的大小超过m，则弹出最大元素，保持堆的大小为m，确保sum_a始终是前k个元素中最小的m个之和。
预处理后缀数组suf_b：
- 从后往前遍历b数组，同样使用大根堆维护从当前元素到末尾的最小的m个元素之和。处理方式与pre_a类似，确保sum_b是当前处理段的最小m个元素之和。
遍历分割点k：
- 遍历所有可能的k值，计算pre_a[k]（前k个元素的最小m个a之和）和suf_b[k+1]（从k+1到末尾的最小m个b之和）的总和，取最小值作为结果。