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
Solution: Priority Queue
Keep all numbers in a priority queue (max heap), each time extract the top one (largest one), then put num – sqrt(num) back to the queue.
Tip: We can early return if all the numbers become 1.
Time complexity: O(n + klogn)
Space complexity: O(n)
C++
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 |
// Author: Huahua class Solution { public: long long pickGifts(vector<int>& gifts, int k) { long long ans = accumulate(begin(gifts), end(gifts), 0LL); priority_queue<int> q(begin(gifts), end(gifts)); while (k-- && q.top() > 1) { int cur = q.top(); q.pop(); int next = sqrt(cur); ans -= (cur - next); q.push(next); } return ans; } }; |