题目描述
给你一个整数数组 gifts
,表示各堆礼物的数量。每一秒,你需要执行以下操作:
- 选择礼物数量最多的那一堆。
- 如果不止一堆都符合礼物数量最多,从中选择任一堆即可。
- 选中的那一堆留下平方根数量的礼物(向下取整),取走其他的礼物。
返回在 k
秒后剩下的礼物数量。
样例
输入:gifts = [25,64,9,4,100], k = 4
输出:29
解释:
按下述方式取走礼物:
- 在第一秒,选中最后一堆,剩下 10 个礼物。
- 接着第二秒选中第二堆礼物,剩下 8 个礼物。
- 然后选中第一堆礼物,剩下 5 个礼物。
- 最后,再次选中最后一堆礼物,剩下 3 个礼物。
最后剩下的礼物数量分别是 [5,8,9,4,3],所以,剩下礼物的总数量是 29。
输入:gifts = [1,1,1,1], k = 4
输出:4
解释:
在本例中,不管选中哪一堆礼物,都必须剩下 1 个礼物。
也就是说,你无法获取任一堆中的礼物。
所以,剩下礼物的总数量是 4。
限制
1 <= gifts.length <= 10^3
1 <= gifts[i] <= 10^9
1 <= k <= 10^3
算法
(堆) $O(n + \min(k, n \log\log m) \log n)$
- 使用堆模拟整个过程,初始化答案为所有礼物数量的总和,模拟过程中逐渐减少。
- 注意到,如果某个礼物的数量变到了 $1$,则无需再次进堆。
时间复杂度
- 建堆的时间复杂度为 $O(n)$。
- 每个数字 $m$ 最多被进堆出堆 $\log\log m$ 次,所以总共会有 $O(\min(k, n \log\log m))$ 次操作堆。
- 故总时间复杂度为 $O(n + \min(k, n \log\log m) \log n)$。
空间复杂度
- 需要 $O(n)$ 的额外空间存储堆。
C++ 代码
#define LL long long
class Solution {
public:
LL pickGifts(vector<int>& gifts, int k) {
priority_queue<int> heap(gifts.begin(), gifts.end());
LL ans = 0;
for (int x : gifts)
ans += x;
while (!heap.empty() && k) {
--k;
int t = heap.top();
heap.pop();
ans -= t - sqrt(t);
t = sqrt(t);
if (t > 1)
heap.push(t);
}
return ans;
}
};