Take Gifts From the Richest Pile

easy heap simulation

Problem

You are given an array gifts where gifts[i] is the number of gifts in pile i. Every second, for k seconds, you walk to the pile that currently has the most gifts and take almost all of them — exactly ⌊√amount⌋ gifts are left behind in that pile. If several piles tie for the most, any one of them may be chosen.

Return the total number of gifts remaining across all piles after k seconds.

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

Output29

100→10, 64→8, 25→5, then the richest is 10→3; piles [5, 8, 9, 4, 3] sum to 29.

gifts (comma-separated)

k (seconds)

import heapq
from math import isqrt

def pick_gifts(gifts, k):
    heap = [-g for g in gifts]   # negate: heapq is a min-heap
    heapq.heapify(heap)
    for _ in range(k):
        largest = -heapq.heappop(heap)
        heapq.heappush(heap, -isqrt(largest))
    return -sum(heap)

function pickGifts(gifts, k) {
  const h = []; // array-backed max-heap
  const push = (v) => {
    h.push(v);
    for (let i = h.length - 1; i > 0; ) {
      const p = (i - 1) >> 1;
      if (h[p] >= h[i]) break;
      [h[p], h[i]] = [h[i], h[p]]; i = p;
    }
  };
  const pop = () => {
    const top = h[0], last = h.pop();
    if (h.length) {
      h[0] = last;
      for (let i = 0; ; ) {
        let m = i; const l = 2 * i + 1, r = l + 1;
        if (l < h.length && h[l] > h[m]) m = l;
        if (r < h.length && h[r] > h[m]) m = r;
        if (m === i) break;
        [h[m], h[i]] = [h[i], h[m]]; i = m;
      }
    }
    return top;
  };
  gifts.forEach(push);
  for (let t = 0; t < k; t++) push(Math.floor(Math.sqrt(pop())));
  return h.reduce((a, b) => a + b, 0);
}

long pickGifts(int[] gifts, int k) {
    PriorityQueue<Integer> heap = new PriorityQueue<>(Collections.reverseOrder());
    for (int g : gifts) heap.offer(g);
    for (int i = 0; i < k; i++) {
        int largest = heap.poll();
        heap.offer((int) Math.sqrt(largest));
    }
    long total = 0;
    for (int g : heap) total += g;
    return total;
}

long long pickGifts(vector<int>& gifts, int k) {
    priority_queue<int> heap(gifts.begin(), gifts.end());
    while (k--) {
        int largest = heap.top();
        heap.pop();
        heap.push((int) sqrt((double) largest));
    }
    long long total = 0;
    while (!heap.empty()) { total += heap.top(); heap.pop(); }
    return total;
}

Explanation

The simulation asks the same question k times in a row: which pile is currently the largest? That is exactly what a max-heap (priority queue) is built for — it hands back the maximum in O(log n) per operation, no matter how the values shuffle around between seconds.

So the whole algorithm is: build a max-heap from gifts once, then repeat k times — pop the largest pile, compute ⌊√largest⌋ (the gifts left behind), and push that value back into the heap. When the loop ends, the answer is simply the sum of everything still in the heap.

Walking through the default example gifts = [25, 64, 9, 4, 100], k = 4: second 1 pops 100 and pushes back ⌊√100⌋ = 10; second 2 pops 64 and leaves 8; second 3 pops 25 and leaves 5; second 4 finds that 10 is now the richest pile and leaves ⌊√10⌋ = 3. The piles end as [5, 8, 9, 4, 3], which sum to 29.

Why not just re-sort the array every second? That works, but each sort costs O(n log n), so the total becomes O(k · n log n). The heap only does local repair work after each pop and push, so every second costs O(log n) — a much better fit when k and n grow.

One detail per language: Python's heapq is a min-heap, so we store negated values and flip the sign on the way in and out; isqrt gives an exact floored integer square root. Java and C++ have native max-ordered priority queues. Also note how fast the square root shrinks piles — once a pile drops to 1, ⌊√1⌋ = 1 keeps it there forever.

Complexity: heapify is O(n), the k seconds cost O(k log n), and the final sum is O(n), giving O(n + k log n) time with O(n) extra space for the heap.

Time: O(n + k log n) Space: O(n)