Take Gifts From the Richest Pile

Easy

#2338Time: O(k * n log n), where `n` is the number of piles and `k` is the number of seconds. In each of the `k` iterations, we sort the array, which takes O(n log n) time. This makes it the slowest approach.Space: O(log n) to O(n). The space complexity depends on the sorting algorithm used. In Java, `Arrays.sort()` for primitive types uses a dual-pivot Quicksort, which requires O(log n) space on average but can degrade to O(n) in the worst case.1 company

Data structures

Array Heap (Priority Queue)

Companies

DE Shaw

Prompt

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
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³
1 <= gifts[i] <= 10⁹
1 <= k <= 10³

Approaches

3 approaches with complexity analysis and trade-offs.

This approach directly simulates the process by sorting the array in each of the k seconds. Sorting makes it easy to find the maximum element, which is always the last element in the sorted array. However, re-sorting the entire array k times is computationally expensive and inefficient.

Algorithm

Repeat the process k times.
In each iteration, sort the entire gifts array in ascending order.
The largest element will be at the end of the sorted array (gifts[n-1]).
Replace this largest element with the floor of its square root.
After k iterations, calculate the sum of all elements in the modified array and return it.

Walkthrough

The most straightforward, yet least efficient, way to solve this problem is to follow the instructions literally but use sorting to simplify finding the maximum value. In each of the k seconds, we sort the gifts array. The pile with the maximum number of gifts will then be the last element. We update this pile's value to the floor of its square root. We repeat this process k times. Finally, we sum up the remaining gifts in the array.

1import java.util.Arrays;2 3class Solution {4    public long pickGifts(int[] gifts, int k) {5        int n = gifts.length;6        for (int i = 0; i < k; i++) {7            // Sort the array to find the max element easily8            Arrays.sort(gifts);9            10            // The last element is the maximum11            int maxGifts = gifts[n - 1];12            13            // Take the gifts and update the pile14            gifts[n - 1] = (int) Math.floor(Math.sqrt(maxGifts));15        }16        17        // Calculate the sum of remaining gifts18        long totalGifts = 0;19        for (int gift : gifts) {20            totalGifts += gift;21        }22        23        return totalGifts;24    }25}

Complexity

Time

O(k * n log n), where `n` is the number of piles and `k` is the number of seconds. In each of the `k` iterations, we sort the array, which takes O(n log n) time. This makes it the slowest approach.

Space

O(log n) to O(n). The space complexity depends on the sorting algorithm used. In Java, `Arrays.sort()` for primitive types uses a dual-pivot Quicksort, which requires O(log n) space on average but can degrade to O(n) in the worst case.

Trade-offs

Pros

Simple to conceptualize as it directly uses a standard library function (sort) to find the maximum.

Cons

Extremely inefficient due to the high cost of sorting the entire array in every single iteration.
The time complexity is significantly worse than other possible solutions.

This is a direct simulation approach. For each of the k seconds, we iterate through the entire gifts array to find the pile with the maximum number of gifts. After finding it, we update its value as per the problem's rule. This is repeated k times. This is more efficient than the sorting approach because finding the maximum in a single pass (O(n)) is faster than sorting (O(n log n)).

Algorithm

Repeat the process k times.
In each iteration, perform a linear scan through the gifts array to find the value and index of the maximum element.
Once the maximum is found, update the element at its index to be the floor of its square root.
After k iterations, sum all the elements in the array to get the final result.

Walkthrough

This method simulates the process step-by-step without any complex data structures. We loop k times. Inside the loop, we initialize variables to track the maximum value found so far and its index. We then iterate through the gifts array from beginning to end. If we find an element larger than our current maximum, we update our tracking variables. After checking all elements, we will have found the index of the largest pile. We then update the value at this index. After the outer loop completes, we compute the sum of all remaining gifts.

1class Solution {2    public long pickGifts(int[] gifts, int k) {3        int n = gifts.length;4        for (int i = 0; i < k; i++) {5            int maxIndex = -1;6            int maxValue = -1;7            8            // Find the pile with the maximum number of gifts9            for (int j = 0; j < n; j++) {10                if (gifts[j] > maxValue) {11                    maxValue = gifts[j];12                    maxIndex = j;13                }14            }15            16            // This check is for cases where all gifts might become 017            if (maxIndex == -1) break;18            19            // Take the gifts by updating the value20            gifts[maxIndex] = (int) Math.floor(Math.sqrt(gifts[maxIndex]));21        }22        23        // Calculate the sum of remaining gifts24        long totalGifts = 0;25        for (int gift : gifts) {26            totalGifts += gift;27        }28        29        return totalGifts;30    }31}

Complexity

Time

O(k * n), where `n` is the number of piles and `k` is the number of seconds. For each of the `k` seconds, we perform a linear scan of the array of size `n` to find the maximum element.

Space

O(1), as we only use a few variables to keep track of the maximum value and its index. The array is modified in-place.

Trade-offs

Pros

Simple to implement and understand.
Very low space complexity as it modifies the input array in-place.

Cons

This approach is inefficient for large k and n values because it repeatedly scans the entire array.
The total number of operations is proportional to the product of k and n.

The most efficient approach utilizes a Max-Heap (Priority Queue). The problem requires repeatedly finding and updating the maximum element in a collection. A max-heap is a data structure specifically designed for this operation, providing logarithmic time complexity for insertions and extractions of the maximum element. This is significantly faster than the linear scan or sorting methods for each step.

Algorithm

Create a Max-Heap data structure. In Java, this can be implemented using a PriorityQueue with a reverse order comparator.
Insert all elements from the gifts array into the max-heap. This process is called heapifying and takes O(n) time.
Loop k times:
- Extract the maximum element from the heap (root of the heap). This is an O(log n) operation.
- Calculate the floor of the square root of the extracted element.
- Insert the new value back into the heap. This is also an O(log n) operation.
After the loop, the heap contains all the final gift counts. Sum up all elements remaining in the heap and return the total.

Walkthrough

We can optimize the process of finding the maximum element by using a max-heap. A max-heap always keeps the largest element at the root, making it accessible in O(1) time (and extractable in O(log n) time).

The algorithm is as follows:

Initialize a PriorityQueue in Java, configured to act as a max-heap.
Add all the initial gift counts to this heap.
Iterate k times. In each iteration, use poll() to remove the largest element, calculate its square root, and use add() to insert the new value back into the heap.
After k seconds, the heap will contain the final amounts of gifts. We can then iterate through the heap (or poll all elements) to sum them up and get the final answer.

1import java.util.PriorityQueue;2import java.util.Collections;3 4class Solution {5    public long pickGifts(int[] gifts, int k) {6        // Create a max-heap using PriorityQueue with a reverse order comparator7        PriorityQueue<Integer> maxHeap = new PriorityQueue<>(Collections.reverseOrder());8        9        // Add all gifts to the max-heap10        for (int gift : gifts) {11            maxHeap.add(gift);12        }13        14        // Perform the operation k times15        for (int i = 0; i < k; i++) {16            // Get the pile with the maximum gifts17            int maxGifts = maxHeap.poll();18            19            // Calculate the remaining gifts and add it back to the heap20            int remainingGifts = (int) Math.floor(Math.sqrt(maxGifts));21            maxHeap.add(remainingGifts);22        }23        24        // Calculate the sum of all gifts remaining in the heap25        long totalGifts = 0;26        for (int gift : maxHeap) {27            totalGifts += gift;28        }29        30        return totalGifts;31    }32}

Complexity

Time

O(n + k log n). Building the heap from the initial array takes O(n) time. Each of the `k` operations involves one extraction (poll) and one insertion (add), both of which take O(log n) time. The final summation takes O(n).

Space

O(n), where `n` is the number of piles. We need to store all `n` gift counts in the priority queue.

Trade-offs

Pros

Highly efficient time complexity, making it suitable for large inputs.
The ideal data structure for problems involving repeated extraction of the minimum or maximum element.

Cons

Requires extra space to store the elements in the heap.
Slightly more complex to implement compared to the brute-force approaches.

Solutions

Solution

1class Solution {2public3  long pickGifts(int[] gifts, int k) {4    PriorityQueue<Integer> pq = new PriorityQueue<>((a, b)->b - a);5    for (int v : gifts) {6      pq.offer(v);7    }8    while (k-- > 0) {9      pq.offer((int)Math.sqrt(pq.poll()));10    }11    long ans = 0;12    for (int v : pq) {13      ans += v;14    }15    return ans;16  }17}

1class Solution {2public:3  long long pickGifts(vector<int> &gifts, int k) {4    make_heap(gifts.begin(), gifts.end());5    while (k--) {6      pop_heap(gifts.begin(), gifts.end());7      gifts.back() = sqrt(gifts.back());8      push_heap(gifts.begin(), gifts.end());9    }10    return accumulate(gifts.begin(), gifts.end(), 0LL);11  }12};

1class Solution:2    def pickGifts(self, gifts: List[int], k: int) -> int: h = [- v for v in gifts] heapify(h) for _ in range(k): heapreplace(h, - int(sqrt(- h[0]))) return - sum(h)

Video walkthrough

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.

Practice

Same difficulty — related problems to reinforce the pattern.

← Previous

Previous problem

Next problem

Solution

1class Solution {2public3  long pickGifts(int[] gifts, int k) {4    PriorityQueue<Integer> pq = new PriorityQueue<>((a, b)->b - a);5    for (int v : gifts) {6      pq.offer(v);7    }8    while (k-- > 0) {9      pq.offer((int)Math.sqrt(pq.poll()));10    }11    long ans = 0;12    for (int v : pq) {13      ans += v;14    }15    return ans;16  }17}

1class Solution {2public:3  long long pickGifts(vector<int> &gifts, int k) {4    make_heap(gifts.begin(), gifts.end());5    while (k--) {6      pop_heap(gifts.begin(), gifts.end());7      gifts.back() = sqrt(gifts.back());8      push_heap(gifts.begin(), gifts.end());9    }10    return accumulate(gifts.begin(), gifts.end(), 0LL);11  }12};

1class Solution:2    def pickGifts(self, gifts: List[int], k: int) -> int: h = [- v for v in gifts] heapify(h) for _ in range(k): heapreplace(h, - int(sqrt(- h[0]))) return - sum(h)

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.