Take Gifts From the Richest Pile

How it works

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.

Steps

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

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.

Things to watch for

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.

from typing import List
from math import sqrt
import heapq

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)

class Solution {
public:
    long long pickGifts(vector<int>& gifts, int k) {
        priority_queue<int> pq(gifts.begin(), gifts.end());

        for (int t = 0; t < k; t++) {
            int n = pq.top(); pq.pop();
            pq.push((int)floor(sqrt(n)));
        }

        long long sum = 0;
        while (!pq.empty()) {
            sum += pq.top(); pq.pop();
        }
        return sum;
    }
};