2558. Take Gifts From the Richest Pile - Explanation

Description

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.
Reduce the number of gifts in the pile to the floor of the square root of the original number of gifts in the pile.

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

Constraints:

1 <= gifts.length <= 1000
1 <= gifts[i] <= 1,000,000,000
1 <= k <= 1000

Company Tags

Please upgrade to NeetCode Pro to view company tags.

1. Simulation

Intuition

The most straightforward approach simulates the process exactly as described. Each second, we find the pile with the most gifts, take gifts from it, and leave behind the floor of its square root. After k seconds, we sum up all remaining gifts. Finding the maximum each time requires scanning all piles.

Algorithm

Repeat k times:
- Find the index of the maximum element in the array.
- Replace that element with the floor of its square root.
Return the sum of all elements.

class Solution:
    def pickGifts(self, gifts: List[int], k: int) -> int:
        for _ in range(k):
            maxIdx = 0
            for i in range(1, len(gifts)):
                if gifts[i] > gifts[maxIdx]:
                    maxIdx = i
            gifts[maxIdx] = int(sqrt(gifts[maxIdx]))
        return sum(gifts)

Time & Space Complexity

Time complexity: $O(n * k)$
Space complexity: $O(1)$ extra space.

Where $n$ is the size of input array, $k$ is the number of seconds.

2. Max-Heap

Intuition

Finding the maximum element repeatedly is expensive with a linear scan. A max-heap keeps the largest element at the top, allowing O(log n) extraction and insertion. Each second, we pop the maximum, compute its square root, and push the result back. This is much faster when k is large relative to n.

Algorithm

Build a max-heap from all gift values.
Repeat k times:
- Pop the maximum value from the heap.
- Push the floor of its square root back into the heap.
Sum all elements remaining in the heap and return the total.

class Solution:
    def pickGifts(self, gifts: List[int], k: int) -> int:
        for i in range(len(gifts)):
            gifts[i] = -gifts[i]
        heapq.heapify(gifts)

        for _ in range(k):
            n = -heapq.heappop(gifts)
            heapq.heappush(gifts, -floor(sqrt(n)))

        return -sum(gifts)

Time & Space Complexity

Time complexity:
- $O(n + k \log n)$ in Python.
- $O(n \log n + k \log n)$ in other languages.
Space complexity: $O(n)$

Where $n$ is the size of input array, $k$ is the number of seconds.

Common Pitfalls

Using a Min-Heap Instead of Max-Heap

This problem requires finding the maximum element repeatedly. Using a min-heap by mistake will give you the smallest pile instead of the richest one. In languages like Python where heapq is a min-heap by default, you need to negate values to simulate a max-heap.

Integer Overflow When Summing

When calculating the final sum of remaining gifts, the result can exceed the range of a 32-bit integer if the input contains large values. Use a 64-bit integer type (like long in Java or int64 in Go) for the sum to avoid overflow.