Take Gifts From the Richest Pile

1. Define the problem

Take Gifts From the Richest Pile

You are given an array gifts, where gifts_i is the number of gifts in pile i, and a number k. Every second, for k seconds, you look at the pile with the most gifts , take gifts away from it until only the square root of its size is left, rounded down, and leave the rest of the piles alone. Return the total number of gifts still on the table after k seconds. A max-heap is the right tool because the only question you ask each second is "which pile is biggest right now?", and the answer changes after every turn.

Constraints

1 ≤ gifts.length ≤ 10³
1 ≤ gifts_i ≤ 10⁹
0 ≤ k ≤ 10³

Example

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

Output29

Explanation Second 1 shrinks 100 to 10, second 2 shrinks 64 to 8, second 3 shrinks 25 to 5 and second 4 shrinks 10 to 3. The piles left are 9, 4, 8, 5 and 3, which add up to 29.

2. Know the words first

In plain terms

Heap: A container that always knows its best item, where best means smallest or largest depending on how you set it up. Adding an item or taking the best item out costs about log n steps, and you never have to sort the whole collection.
Backing array: A heap is stored as one plain list. The item at position i keeps its parent at position (i - 1) / 2 rounded down, and its two children at positions 2i + 1 and 2i + 2. That is why every picture below is a row of numbered boxes.
Max-heap: A heap whose best item is the largest one. Popping it always hands you the biggest value still inside.
Square root rounded down: Math.floor(Math.sqrt(x)) in code. For 100 it is 10, for 64 it is 8, and for 10 it is 3, because 3 times 3 is 9 and 4 times 4 would be too big.

3. Visualize the solution

The row of boxes is the backing array of the max-heap for gifts = [25, 64, 9, 4, 100]

0100

164

425

Statusbuild

Push all five piles. Slot 0 holds the biggest pile, 100.

What happens in this step

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

Slot 0 holds 100 and its children in slots 1 and 2 hold 64 and 9.
No child is bigger than its parent, so this is a valid max-heap.

Step 1 of 6

Steps to visualize

Each box is one slot of the list that stores the heap. Slot 0 always holds the biggest pile.
Each second, pop slot 0, shrink that number to its square root rounded down, and push it back.
The heap never changes size here, because one pile comes out and one pile goes back in.
After k seconds, add up whatever is left in the heap.

4. Walk through the code

Walk through the code

Same walkthrough, now with the code. Press Next to move one step and watch which lines run.

The row of boxes is the backing array of the max-heap for gifts = [25, 64, 9, 4, 100]

0100

164

425

Statusbuild

Push all five piles. Slot 0 holds the biggest pile, 100.

What happens in this step

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

Slot 0 holds 100 and its children in slots 1 and 2 hold 64 and 9.
No child is bigger than its parent, so this is a valid max-heap.

Step 1 of 6

5. Solution

Solution

solution.tsTypeScript

function pickGifts(gifts, k) {
  const heap = new Heap((a, b) => b - a);

  for (const gift of gifts) {
    heap.push(gift);
  }

  for (let turn = 0; turn < k; turn++) {
    const richest = heap.pop();
    heap.push(Math.floor(Math.sqrt(richest)));
  }

  let total = 0;

  while (heap.size() > 0) {
    total += heap.pop();
  }

  return total;
}

class Heap {
  constructor(compare) {
    this.items = [];
    this.compare = compare;
  }

  size() {
    return this.items.length;
  }

  peek() {
    return this.items[0];
  }

  push(value) {
    this.items.push(value);
    let child = this.items.length - 1;

    while (child > 0) {
      const parent = (child - 1) >> 1;
      if (this.compare(this.items[child], this.items[parent]) >= 0) break;
      const swap = this.items[child];
      this.items[child] = this.items[parent];
      this.items[parent] = swap;
      child = parent;
    }
  }

  pop() {
    const top = this.items[0];
    const last = this.items.pop();

    if (this.items.length > 0) {
      this.items[0] = last;
      let parent = 0;

      while (true) {
        const left = parent * 2 + 1;
        const right = parent * 2 + 2;
        let best = parent;

        if (left < this.items.length && this.compare(this.items[left], this.items[best]) < 0) best = left;
        if (right < this.items.length && this.compare(this.items[right], this.items[best]) < 0) best = right;
        if (best === parent) break;

        const swap = this.items[parent];
        this.items[parent] = this.items[best];
        this.items[best] = swap;
        parent = best;
      }
    }

    return top;
  }
}

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

6. Test cases

Test cases

Input	Expected	Covers
`gifts = [25, 64, 9, 4, 100], k = 4`	`29`	example from the docstring
`gifts = [1, 1, 1, 1], k = 4`	`4`	piles of 1 never shrink because the square root of 1 is 1
`gifts = [9], k = 1`	`3`	smallest input, one pile shrunk once
`gifts = [100], k = 3`	`1`	the same pile is chosen every second, 100 to 10 to 3 to 1
`gifts = [16, 16], k = 2`	`8`	two piles of equal size
`gifts = [3, 2, 1], k = 0`	`6`	no seconds pass, so nothing changes

Take Gifts From the Richest Pile