The gift pile reduction

You are given piles of gifts. Each second, you pick the pile with the most gifts, leave behind the whole-number square root of it, and take the rest. After a set number of seconds, you want to know how many gifts remain in total.

Picture a generous neighbour who keeps raiding the biggest pile in the shop.

Always the largest: each second, find the current biggest pile.

Square root: the pile shrinks to its whole-number square root.

gifts = [25, 64, 9, 4, 100], k = 4

biggest

total: 112

Second 1: the biggest pile is 100. Leave √100 = 10.

Move 1 of 4

Efficiency with a max-heap

Because you only change one pile each second, the other piles stay where they were. Pushing and popping on a max-heap fixes the order in logarithmic time, so you don't re-sort everything each time.

Do this for the given number of seconds, then add up what is left.

Fast enough: k operations, each about log n steps.

64-bit integer: the sum of all remaining piles can exceed a 32-bit int.

In code

import heapq
import math

heap = [-g for g in gifts]
heapq.heapify(heap)
for _ in range(k):
    m = -heapq.heappop(heap)
    heapq.heappush(heap, -math.isqrt(m))
return -sum(heap)