给你一个整数数组 gifts
,表示各堆礼物的数量。每一秒,你需要执行以下操作:
- 选择礼物数量最多的那一堆。
- 如果不止一堆都符合礼物数量最多,从中选择任一堆即可。
- 选中的那一堆留下平方根数量的礼物(向下取整),取走其他的礼物。
返回在 k
秒后剩下的礼物数量。
示例 1:
输入:gifts = [25,64,9,4,100], k = 4 输出:29 解释: 按下述方式取走礼物: - 在第一秒,选中最后一堆,剩下 10 个礼物。 - 接着第二秒选中第二堆礼物,剩下 8 个礼物。 - 然后选中第一堆礼物,剩下 5 个礼物。 - 最后,再次选中最后一堆礼物,剩下 3 个礼物。 最后剩下的礼物数量分别是 [5,8,9,4,3] ,所以,剩下礼物的总数量是 29 。
示例 2:
输入:gifts = [1,1,1,1], k = 4 输出:4 解释: 在本例中,不管选中哪一堆礼物,都必须剩下 1 个礼物。 也就是说,你无法获取任一堆中的礼物。 所以,剩下礼物的总数量是 4 。
提示:
1 <= gifts.length <= 10^3
1 <= gifts[i] <= 10^9
1 <= k <= 10^3
提示 1
How can you keep track of the largest gifts in the array
提示 2
What is an efficient way to find the square root of a number?
提示 3
Can you keep adding up the values of the gifts while ensuring they are in a certain order?
提示 4
Can we use a priority queue or heap here?
解法1:最大堆(优先队列)
我们可以用最大堆来维护各堆礼物的数量,进行 k 次如下操作:每次从中取出最大的一个,求根之后再将结果放入堆中。
最后,最大堆中所有礼物的数量之和就是我们要返回的答案。
Java版:
class Solution {
public long pickGifts(int[] gifts, int k) {
PriorityQueue<Integer> pq = new PriorityQueue<>((a, b) -> b - a);
for (int g : gifts) {
pq.offer(g);
}
for (int i = 0; i < k; i++) {
pq.offer((int) Math.sqrt(pq.poll()));
}
long ans = 0;
while (!pq.isEmpty()) {
ans += pq.poll();
}
return ans;
}
}
Python3版:
class Solution:
def pickGifts(self, gifts: List[int], k: int) -> int:
heap = [-g for g in gifts]
heapify(heap)
while k > 0:
x = heappop(heap)
heappush(heap, -int(math.sqrt(-x)))
k -= 1
return -sum(heap)
复杂度分析
- 时间复杂度:O(klogn),其中 n 为 gifts 的长度。初始化建堆复杂度最低可以为 O(n),求解平方根视作 O(1) 的复杂度,元素出堆和入堆的时间复杂度为 O(logn),因此总体复杂度为 O(klogn)。如果使用依次将 n 个元素放入堆中的方式来进行初始化,则时间复杂度为 O((n+k)logn)。
- 空间复杂度:O(n),其中 n 为 gifts 的长度。使用优先队列的空间复杂度为 O(n)。