Python 优先队列：heapq库的使用

小嗷犬2023-03-152023-08-25

简介

heapq 库是 Python 标准库中的一部分，它提供了一些堆操作的函数，可以用来实现优先队列。

优先队列是一种特殊的队列，它的每个元素都有一个优先级，元素的出队顺序是按照优先级从高到低的顺序进行的。优先队列的实现有多种方式，其中最常用的是堆。

堆是一种特殊的树，有两种类型，分别是最大堆和最小堆。最大堆的每个节点的值都大于或等于其子节点的值，最小堆的每个节点的值都小于或等于其子节点的值。堆的根节点是堆中的最大值（最小堆的根节点是最小值）。

heapq 的大部分操作都是基于最小堆实现的，通过将元素取相反数，可以实现最大堆。

heapq 库的使用

heapq 库提供了 heapify、heappush、heappop、heapreplace、heappushpop、merge、nlargest、nsmallest 等函数，用于堆的操作。

heapify

heapify 函数用于原地将列表转换为最小堆，时间复杂度为 $O(n)$ 。

函数原型如下：

1	`heapq.heapify(x)`

其中，x 是一个列表。

示例：

import heapq
x = [1, 3, 5, 7, 9, 2, 4, 6, 8, 0]
heapq.heapify(x)
print(x)
# [0, 1, 2, 6, 3, 5, 4, 7, 8, 9]

heappush

heappush 函数用于将元素插入到最小堆中，并保持堆的不变性，时间复杂度为 $O(\log n)$ 。

函数原型如下：

1	`heapq.heappush(heap, item)`

其中，heap 是一个最小堆，item 是要插入的元素。

示例：

import heapq
x = [1, 3, 5, 7, 9, 2, 4, 6, 8, 0]
heapq.heapify(x)
heapq.heappush(x, 2.5)
print(x)
# [0, 1, 2, 6, 2.5, 5, 4, 7, 8, 9, 3]

heappop

heappop 函数用于弹出最小堆的根节点，并保持堆的不变性，时间复杂度为 $O(\log n)$ 。如果堆为空，则抛出 IndexError 异常。

函数原型如下：

1	`heapq.heappop(heap)`

其中，heap 是一个最小堆。

示例：

import heapq
x = [1, 3, 5, 7, 9, 2, 4, 6, 8, 0]
heapq.heapify(x)
print(heapq.heappop(x))
# 0
print(x)
# [1, 3, 2, 6, 9, 5, 4, 7, 8]

使用 heap[0] 可以访问最小堆的根节点，但是不会弹出它。

import heapq
x = [1, 3, 5, 7, 9, 2, 4, 6, 8, 0]
heapq.heapify(x)
print(x[0])
# 0
print(x)
# [0, 1, 2, 6, 3, 5, 4, 7, 8, 9]

heapreplace

heapreplace 函数用于弹出最小堆的根节点，并将新元素插入到堆中，保持堆的大小和不变性，时间复杂度为 $O(\log n)$ 。如果堆为空，则抛出 IndexError 异常。它比先调用 heappop 再调用 heappush 效率更高。

函数原型如下：

1	`heapq.heapreplace(heap, item)`

其中，heap 是一个最小堆，item 是要插入的元素。

示例：

import heapq
x = [1, 3, 5, 7, 9, 2, 4, 6, 8, 0]
heapq.heapify(x)
heapq.heapreplace(x, -1)
print(x)
# [-1, 1, 2, 6, 3, 5, 4, 7, 8, 9]

heappushpop

heappushpop 函数用于将元素插入到最小堆中，并弹出最小堆的根节点，保持堆的大小和不变性，时间复杂度为 $O(\log n)$ 。如果堆为空，则抛出 IndexError 异常。它比先调用 heappush 再调用 heappop 效率更高。

函数原型如下：

1	`heapq.heappushpop(heap, item)`

其中，heap 是一个最小堆，item 是要插入的元素。

示例：

import heapq
x = [1, 3, 5, 7, 9, 2, 4, 6, 8, 0]
heapq.heapify(x)
heapq.heappushpop(x, -1)
print(x)
# [0, 1, 2, 6, 3, 5, 4, 7, 8, 9]

merge

merge 函数是一个基于堆的通用功能函数，用于合并多个有序的序列，返回一个新的有序的序列，时间复杂度为 $O(n \log k)$ ，其中 $n$ 是所有序列的元素个数， $k$ 是序列的个数。函数返回一个已排序值的迭代器，可以使用 list 函数将其转换为列表。

函数原型如下：

1	`heapq.merge(*iterables, key=None, reverse=False)`

其中，iterables 是多个有序的序列，key 是一个函数，用于从序列中提取比较的键，reverse 是一个布尔值，表示是否反转序列。

示例：

import heapq
x = [1, 3, 5, 7, 9]
y = [2, 4, 6, 8, 10]
z = heapq.merge(x, y)
print(list(z))
# [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

nlargest

nlargest 函数是一个基于堆的通用功能函数，用于返回最大的 $n$ 个元素，时间复杂度为 $O(n \log k)$ ，其中 $n$ 是序列的长度， $k$ 是要返回的元素个数。如果 $n$ 小于 $k$ ，则返回整个序列。

函数原型如下：

1	`heapq.nlargest(n, iterable, key=None)`

其中，n 是要返回的元素个数，iterable 是一个序列，key 是一个函数，用于从序列中提取比较的键。

示例：

import heapq
x = [1, 3, 5, 7, 9, 2, 4, 6, 8, 0]
print(heapq.nlargest(3, x))
# [9, 8, 7]

nlargest 函数在 $n$ 值较小时性能较好。对于较大的 $n$ ，使用 sorted(iterable, reverse=True)[:n] 性能更好。当 $n=1$ 时，使用 max(iterable) 函数性能更好。

nsmallest

nsmallest 函数是一个基于堆的通用功能函数，用于返回最小的 $n$ 个元素，时间复杂度为 $O(n \log k)$ ，其中 $n$ 是序列的长度， $k$ 是要返回的元素个数。如果 $n$ 小于 $k$ ，则返回整个序列。

函数原型如下：

1	`heapq.nsmallest(n, iterable, key=None)`

其中，n 是要返回的元素个数，iterable 是一个序列，key 是一个函数，用于从序列中提取比较的键。

示例：

import heapq
x = [1, 3, 5, 7, 9, 2, 4, 6, 8, 0]
print(heapq.nsmallest(3, x))
# [0, 1, 2]

nsmallest 函数在 $n$ 值较小时性能较好。对于较大的 $n$ ，使用 sorted(iterable)[:n] 性能更好。当 $n=1$ 时，使用 min(iterable) 函数性能更好。

例题

Title

CodeForces 1800 C2. Powering the Hero (hard version)

Time Limit

2 seconds

Memory Limit

256 megabytes

Problem Description

This is a hard version of the problem. It differs from the easy one only by constraints on $n$ and $t$ .

There is a deck of $n$ cards, each of which is characterized by its power. There are two types of cards:

You can do the following with the deck:

Your task is to use such actions to gather an army with the maximum possible total power.

Input

The first line of input data contains single integer $t$ ( $1 \le t \le 10^4$ ) — the number of test cases in the test.

The first line of each test case contains one integer $n$ ( $1 \le n \le 2 \cdot 10^5$ ) — the number of cards in the deck.

The second line of each test case contains $n$ integers $s_1, s_2, \dots, s_n$ ( $0 \le s_i \le 10^9$ ) — card powers in top-down order.

It is guaranteed that the sum of $n$ over all test cases does not exceed $2 \cdot 10^5$ .

Output

Output $t$ numbers, each of which is the answer to the corresponding test case — the maximum possible total power of the army that can be achieved.

Sample Input

5
5
3 3 3 0 0
6
0 3 3 0 0 3
7
1 2 3 0 4 5 0
7
1 2 5 0 4 3 0
5
3 1 0 0 4

Sample Onput

Note

In the first sample, you can take bonuses $1$ and $2$ . Both hero cards will receive $3$ power. If you take all the bonuses, one of them will remain unused.

In the second sample, the hero’s card on top of the deck cannot be powered up, and the rest can be powered up with $2$ and $3$ bonuses and get $6$ total power.

In the fourth sample, you can take bonuses $1$ , $2$ , $3$ , $5$ and skip the bonus $6$ , then the hero $4$ will be enhanced with a bonus $3$ by $5$ , and the hero $7$ with a bonus $5$ by $4$ . $4+5=9$ .

Source

CodeForces 1800 C2. Powering the Hero (hard version)

Solution

每张英雄牌的最大力量为该英雄牌之前出现的未被使用最大奖励牌的力量。对于具体是哪张英雄牌使用了哪张奖励牌，我们是不关心的，只需要统计他们最大力量即可。

import heapq

for _ in range(int(input())):
    n = int(input())
    s = map(int, input().split())
    h = []
    ans = 0
    for i in s:
        if i == 0 and h:
            ans -= heapq.heappop(h)
        else:
            heapq.heappush(h, -i)
    print(ans)