LC 2558. 从数量最多的堆取走礼物
题目描述
这是 LeetCode 上的 2558. 从数量最多的堆取走礼物 ，难度为 简单。
给你一个整数数组 gifts，表示各堆礼物的数量。
每一秒，你需要执行以下操作：
选择礼物数量最多的那一堆。
如果不止一堆都符合礼物数量最多，从中选择任一堆即可。
选中的那一堆留下平方根数量的礼物（向下取整），取走其他的礼物。
返回在 k 秒后剩下的礼物数量。
示例 1：1
2
3
4
5
6
7
8
9
10
11输入：gifts = [25,64,9,4,100], k = 4
输出：29
解释：
按下述方式取走礼物：
- 在第一秒，选中最后一堆，剩下 10 个礼物。
- 接着第二秒选中第二堆礼物，剩下 8 个礼物。
- 然后选中第一堆礼物，剩下 5 个礼物。
- 最后，再次选中最后一堆礼物，剩下 3 个礼物。
最后剩下的礼物数量分别是 [5,8,9,4,3] ，所以，剩下礼物的总数量是 29 。
示例 2：1
2
3
4
5
6
7
8输入：gifts = [1,1,1,1], k = 4
输出：4
解释：
在本例中，不管选中哪一堆礼物，都必须剩下 1 个礼物。
也就是说，你无法获取任一堆中的礼物。
所以，剩下礼物的总数量是 4 。
提示：
- $1 <= gifts.length <= 10^3$
- $1 <= gifts[i] <= 10^9$
- $1 <= k <= 10^3$
基本思想
为了方便，将 gifts 记为 gs。
从题面看，数组大小和执行次数范围均为 $1e3$，那么暴力做法的复杂度为 $O(n^2)$，计算量为 $1e6$，是可以通过的。
稍有数据结构基础的同学，容易进一步联想：存在一种数据结构，能够快速查得集合最值，使复杂度降低。
该数据结构是「堆」，在某些 STL 里面又称为「优先队列」。
手写堆入门
所有高级数据结构（例如之前讲过的 [字典树 Trie] 或 [并查集]，都可以使用数组进行模拟，堆自然也不例外。
堆本质是数据集合，只不过区别于简单数组而言，其具有 $O(1)$ 查找最值特性。
既然我们使用数组来实现堆，同时又需要实现 $O(1)$ 查找最值，一个自然而然的想法，是把集合最值放到数组的特定位置，例如数组起始位置。
开始之前，先来了解一下，堆的基本操作：
add：往堆中增加元素peek：快速查得最值pop：将最值从堆中移除
我们以小根堆为例，进行学习。
PS. 大根堆亦是同理，如果已经有写好的小根堆模板，那么将数值进行符号翻转，即可实现大根堆；或是翻转模板中的数值比较逻辑，也可将小根堆轻松切换成大根堆。
1. 堆长啥样？
堆是数组实现的，但其形态为「完全二叉树」（最多只有底下一层节点不满），目的是尽量让树平衡，降低树的高度，从而更高效的维护其性质。
虽是「完全二叉树」，但其又和另外一种特殊二叉树，二叉搜索树（BST）不同。
在 BST 的定义中：任意节点的左子树上的节点值均不大于该节点，任意节点的右子树上的节点值均不小于该节点。
在小根堆对应的「完全二叉树」定义中：任意节点的节点值均不大于其左右子树中的节点值。
因此将小根堆可视化后，也十分形象：小根堆的堆顶元素最小，然后从上往下，元素“依次”增大。
{:width=400}
即对于小根堆而言，每个子树的根节点，都是该子树的最小值。
2. 如何用数组进行「完全二叉树」的存储？
我们只需要将「二叉树的左右子节点关系」和「数组下标」实现某种关联即可。
我们约定，对于某个节点，若其所在数组下标为 $idx$，那么该节点的左子节点的数组下标为 $2 \times idx$，该节点的右子节点的数组下标为 $2 \times idx + 1$。
注意：基于此规则，我们需要调整数组下标从 $1$ 开始进行存储。
{:width=400}
3.基本操作如何实现？
所有操作的核心最终都归结到：如何通过有限的调整，重新恢复堆所对应的完全二叉树中的节点定义。
因此可抽象出 up 操作和 down 操作，俩操作均通过递归实现：
up操作是将当前节点u与父节点fa = u / 2进行比较，若发现父节点更大，则将两者进行互换，并递归处理父节点down操作是将当前节点cur和左右节点l = cur * 2和r = cur * 2 + 1进行比较，若当前节点并非三者最小，则进行互换，并递归处理子节点
通过上述核心 up 和 down 操作，以及「完全二叉树」在数组的放置规则，我们可以容易拼凑出堆的基本操作。
假定当前使用一维数组 heap 实现堆，使用变量 sz 记录当前堆的元素个数，同时该变量充当 heap 的结尾游标：
add：往堆中增加元素
该操作可转换为：往数组尾部增加元素（heap[sz++] = x），并从尾部开始重整堆（up(sz)）peek：快速查得最小值
直接范围数组首位元素，注意下标从 $1$ 开始，heap[1]pop：将最值从堆中移除
先通过peek操作记录下待移除的最小值，然后将数组的结尾元素和首位元素互换，并更新数组长度减一（heap[1] = heap[sz--]），随后从头部开始重整堆（down(1)）
Java 完整模板：1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39int[] heap = new int[10010];
int sz = 0;
void swap(int a, int b) {
int c = heap[a];
heap[a] = heap[b];
heap[b] = c;
}
void up(int u) {
// 将「当前节点 i」与「父节点 i / 2」进行比较, 若父节点值更大, 则进行交换
int fa = u / 2;
if (fa != 0 && heap[fa] > heap[u]) {
swap(fa, u);
up(fa);
}
}
void down(int u) {
// 将当「前节点 cur」与「左节点 l」及「右节点 r」进行比较, 找出最小值, 若当前节点不是最小值, 则进行交换
int cur = u;
int l = u * 2, r = u * 2 + 1;
if (l <= sz && heap[l] < heap[cur]) cur = l;
if (r <= sz && heap[r] < heap[cur]) cur = r;
if (cur != u) {
swap(cur, u);
down(cur);
}
}
void add(int x) {
heap[++sz] = x;
up(sz);
}
int peek() {
return heap[1];
}
int pop() {
int ans = peek();
heap[1] = heap[sz--];
down(1);
return ans;
}
C++ 完整模板：1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32int heap[1010];
int sz = 0;
void up(int u) {
int fa = u / 2;
if (fa && heap[fa] >= heap[u]) {
swap(heap[fa], heap[u]);
up(fa);
}
}
void down(int u) {
int cur = u;
int l = cur * 2, r = cur * 2 + 1;
if (l <= sz && heap[l] < heap[cur]) cur = l;
if (r <= sz && heap[r] < heap[cur]) cur = r;
if (cur != u) {
swap(heap[cur], heap[u]);
down(cur);
}
}
void add(int x) {
heap[++sz] = x;
up(sz);
}
int peek() {
return heap[1];
}
int poll() {
int ans = peek();
heap[1] = heap[sz--];
down(1);
return ans;
}
模板运用
利用「上述模板」以及「小根堆如何当大根堆」使用，我们可以轻松解决本题。
- 起始先将逐元素进行符号翻转并入堆
- 执行
k次取出并重放操作（注意符号转换） - 统计所有元素之和，由于此时不再需要查询最值，可通过直接扫描数组的方式（注意符号转换）
Java 代码：1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48class Solution {
int[] heap = new int[10010];
int sz = 0;
public long pickGifts(int[] gs, int k) {
for (int x : gs) add(-x);
while (k-- > 0) add(-(int)Math.sqrt(-pop()));
long ans = 0;
while (sz != 0) ans += -heap[sz--]; // 没必要再维持堆的有序, 直接读取累加
return ans;
}
void swap(int a, int b) {
int c = heap[a];
heap[a] = heap[b];
heap[b] = c;
}
void up(int u) {
// 将「当前节点 i」与「父节点 i / 2」进行比较, 若父节点值更大, 则进行交换
int fa = u / 2;
if (fa != 0 && heap[fa] > heap[u]) {
swap(fa, u);
up(fa);
}
}
void down(int u) {
// 将当「前节点 cur」与「左节点 l」及「右节点 r」进行比较, 找出最小值, 若当前节点不是最小值, 则进行交换
int cur = u;
int l = u * 2, r = u * 2 + 1;
if (l <= sz && heap[l] < heap[cur]) cur = l;
if (r <= sz && heap[r] < heap[cur]) cur = r;
if (cur != u) {
swap(cur, u);
down(cur);
}
}
void add(int x) {
heap[++sz] = x;
up(sz);
}
int peek() {
return heap[1];
}
int pop() {
int ans = peek();
heap[1] = heap[sz--];
down(1);
return ans;
}
}
C++ 代码：1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42class Solution {
public:
int heap[1010];
int sz = 0;
long long pickGifts(vector<int>& gs, int k) {
for (int x : gs) add(-x);
while (k--) add(-sqrt(-poll()));
long ans = 0;
while (sz != 0) ans += -heap[sz--];
return ans;
}
void up(int u) {
int fa = u / 2;
if (fa && heap[fa] >= heap[u]) {
swap(heap[fa], heap[u]);
up(fa);
}
}
void down(int u) {
int cur = u;
int l = cur * 2, r = cur * 2 + 1;
if (l <= sz && heap[l] < heap[cur]) cur = l;
if (r <= sz && heap[r] < heap[cur]) cur = r;
if (cur != u) {
swap(heap[cur], heap[u]);
down(cur);
}
}
void add(int x) {
heap[++sz] = x;
up(sz);
}
int peek() {
return heap[1];
}
int poll() {
int ans = peek();
heap[1] = heap[sz--];
down(1);
return ans;
}
};
- 时间复杂度：建堆复杂度为 $O(n\log{n})$；执行
k次操作复杂度为 $O(k\log{n})$；构建答案复杂度为 $O(n)$。整体复杂度为 $O((n + k) \log{n})$ - 空间复杂度：$O(n)$
最后
这是我们「刷穿 LeetCode」系列文章的第 No.2558 篇，系列开始于 2021/01/01，截止于起始日 LeetCode 上共有 1916 道题目，部分是有锁题，我们将先把所有不带锁的题目刷完。
在这个系列文章里面，除了讲解解题思路以外，还会尽可能给出最为简洁的代码。如果涉及通解还会相应的代码模板。
为了方便各位同学能够电脑上进行调试和提交代码，我建立了相关的仓库：https://github.com/SharingSource/LogicStack-LeetCode 。
在仓库地址里，你可以看到系列文章的题解链接、系列文章的相应代码、LeetCode 原题链接和其他优选题解。
本博客所有文章除特别声明外，均采用 CC BY-SA 4.0 协议 ，转载请注明出处！
