Take Gifts From the Richest Pile

lcpy gen -n 2558   # by problem number
lcpy gen -s take_gifts_from_the_richest_pile   # by problem name

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.

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.

import heapq
from math import isqrt


class Solution:
    # Time: O(k log n + n)
    # Space: O(n)
    def pick_gifts(self, gifts: list[int], k: int) -> int:
        heap = [-gift for gift in gifts]
        heapq.heapify(heap)
        for _ in range(k):
            top = -heapq.heappop(heap)
            if top <= 1:
                heapq.heappush(heap, -top)
                break
            heapq.heappush(heap, -isqrt(top))
        return -sum(heap)