[LeetCode] 2558. Take Gifts From the Richest Pile

You are given an integer array gifts denoting the number of gifts in various piles. Every second, you do the following:

Choose the pile with the maximum number of gifts.
If there is more than one pile with the maximum number of gifts, choose any.
Leave behind the floor of the square root of the number of gifts in the pile. Take the rest of the gifts.
Return the number of gifts remaining after k seconds.

Example 1:
Input: gifts = [25,64,9,4,100], k = 4
Output: 29
Explanation:
The gifts are taken in the following way:

• In the first second, the last pile is chosen and 10 gifts are left behind.
• Then the second pile is chosen and 8 gifts are left behind.
• After that the first pile is chosen and 5 gifts are left behind.
• Finally, the last pile is chosen again and 3 gifts are left behind.
  The final remaining gifts are [5,8,9,4,3], so the total number of gifts remaining is 29.

Example 2:
Input: gifts = [1,1,1,1], k = 4
Output: 4
Explanation:
In this case, regardless which pile you choose, you have to leave behind 1 gift in each pile.
That is, you can’t take any pile with you.
So, the total gifts remaining are 4.

Constraints:
1 <= gifts.length <= 103
1 <= gifts[i] <= 109
1 <= k <= 103

Follow-up: Can you come up with an algorithm that is less than O(n2) time complexity?

从数量最多的堆取走礼物。

给你一个整数数组 gifts ，表示各堆礼物的数量。每一秒，你需要执行以下操作：
选择礼物数量最多的那一堆。
如果不止一堆都符合礼物数量最多，从中选择任一堆即可。
选中的那一堆留下平方根数量的礼物（向下取整），取走其他的礼物。
返回在 k 秒后剩下的礼物数量。

思路

思路是最大堆。首先创建一个最大堆将每一堆礼物放进去，放进去的同时累加一下礼物的个数，记为 res。
接着开始从最大堆中弹出元素 top。弹出堆顶元素的时候记得要从 res 中减去 top 的平方根，然后将剩余的部分再放回最大堆，同时记得 k–。
注意如果遇到一堆礼物的个数 = 1的时候，说明这一堆礼物无法再分割了，此时即可跳出 while 循环返回 res。

复杂度

时间O(nlogn)
空间O(n)

代码

Java实现

class Solution {
    public long pickGifts(int[] gifts, int k) {
        long res = 0;
        PriorityQueue<Integer> queue = new PriorityQueue<>((a, b) -> b - a);
        for (int g : gifts) {
            queue.offer(g);
            res += g;
        }

        int kk = k;
        while (kk > 0 && !queue.isEmpty()) {
            int top = queue.poll();
            if (top == 1) {
                break;
            }
            int taken = top - (int) Math.sqrt(top);
            res -= taken;
            int rest = top - taken;
            queue.offer(rest);
            kk--;
        }
        return res;
    }
}