Problem

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 afterk seconds.

Examples

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 <= 10^3
  • 1 <= gifts[i] <= 10^9
  • 1 <= k <= 10^3

Solution

Method 1 - Using Max Heap

To solve this problem, we need to repeatedly choose the pile with the maximum number of gifts, and then update the pile according to the described process. The efficient way to do this is to use a max-heap data structure since it allows us to efficiently retrieve and update the maximum element.

Steps

  1. Convert all elements in the gifts array to negative and build a max-heap (using Python’s heapq which is a min-heap by default).
  2. Perform the operation k times:
    • Extract the maximum element from the heap.
    • Calculate the floor of the square root of this element.
    • Negate and push the updated value back into the heap.
  3. The sum of all elements left in the heap after k seconds is the desired result.

Video explanation

Here is the video explaining this method in detail. Please check it out:

Code

Java
class Solution {
    public long pickGifts(int[] gifts, int k) {
        PriorityQueue<Integer> maxHeap = new PriorityQueue<>((a, b) -> b - a); // MAX-HEAP
        
        // Add all gift counts to the max-heap
        for (int gift : gifts) {
            maxHeap.add(gift);
        }
        
        // Perform the operation k times
        for (int i = 0; i < k; i++) {
            int maxGift = maxHeap.poll();
            int remainingGift = (int) Math.sqrt(maxGift);
            maxHeap.add(remainingGift);
        }
        
        // Sum all remaining gifts in the heap
        long sum = 0;
        while (!maxHeap.isEmpty()) {
            sum += maxHeap.poll();
        }
        
        return sum;
    }
}
Python
class Solution:
    def pick_gifts(self, gifts: List[int], k: int) -> int:
        max_heap = []
        for gift in gifts:
            heappush(max_heap, -gift)  # We use negative values to simulate a max-heap
        
        for _ in range(k):
            max_gift = -heappop(max_heap)
            remaining_gift = math.isqrt(max_gift)
            heappush(max_heap, -remaining_gift)
        
        ans = -sum(max_heap)
        return int(ans)

Complexity

  • ⏰ Time complexity: O(n + k log n), where n is length of the array. O(n) for building the heap, and it takes O(log n) for heap push and pop operations, and we do it k times, hence O(k log(n)), so total time complexity is O(n + k log n).
  • 🧺 Space complexity: O(n) to store all elements of array in heap