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

堆的定义

堆（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 模块只提供最小堆的功能。

【模板】堆

题目描述

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

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

输入格式

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

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

输出格式

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

输入输出样例 #1

输入 #1

5
1 2
1 5
2
3
2

输出 #1

2
5

说明/提示

【数据规模与约定】

• 对于 $30\%$ 的数据，保证 $n\le 15$。
• 对于 $70\%$ 的数据，保证 $n\le 10^4$。
• 对于 $100\%$ 的数据，保证 $1\le n\le 10^6$，$1\le x<2^{31}$， o p ∈ { 1 , 2 , 3 } op \in \{1, 2, 3\} op∈{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题 和+和

方法思路

1. 预处理前缀和后缀数组：

   • pre_a[k]：表示在前k个元素中选择m个最小的a数组元素之和。
   • suf_b[k]：表示从第k个元素到末尾中选择m个最小的b数组元素之和。

2. 使用大根堆维护最小元素：

   • 遍历数组时维护一个大根堆，确保堆中始终保存当前最小的m个元素。当堆的大小超过m时，弹出最大的元素，保持堆的大小为m。

3. 遍历所有可能的分割点：

   • 对于每个可能的分割点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)

代码解释

1. 预处理前缀数组pre_a：

   • 使用大根堆维护前k个元素中最小的m个元素之和。每次添加元素后，若堆的大小超过m，则弹出最大元素，保持堆的大小为m，确保sum_a始终是前k个元素中最小的m个之和。

2. 预处理后缀数组suf_b：

   • 从后往前遍历b数组，同样使用大根堆维护从当前元素到末尾的最小的m个元素之和。处理方式与pre_a类似，确保sum_b是当前处理段的最小m个元素之和。

3. 遍历分割点k：

   • 遍历所有可能的k值，计算pre_a[k]（前k个元素的最小m个a之和）和suf_b[k+1]（从k+1到末尾的最小m个b之和）的总和，取最小值作为结果。

这种方法利用堆高效地维护了最小的m个元素之和，预处理时间复杂度为O(n log m)，遍历分割点的复杂度为O(n)，整体复杂度为O(n log m)，适用于大规模数据。

END
如果有更多问题或需要进一步的帮助，可以在评论区留言讨论哦！
如果喜欢的话，请给博主点个关注 谢谢