Last Stone Weight

The problem

You have a pile of stones, and stones lists their weights, which are positive whole numbers. Smash the two heaviest stones together, and repeat for as long as at least two stones are left. If the two weigh the same, both are destroyed. If not, the lighter one is destroyed and the heavier one loses the lighter one's weight. Return the weight of the last stone, or 0 if no stone is left. The list may be empty.

What to write

function solve(stones: number[]): number

Input
stones = [2,7,4,1,8,1]

Output
1

Input
stones = [1]

Output
1

Constraints

0–100 positive integer weights, each ≤ 10,000.

Hints

Each hint gives away a little more. Open one only when you need it.

1. Hint 1: Use the constraints

   Every round destroys at least one stone, so there are never more rounds than stones. The rounds are not the costly part. What costs time is finding the two heaviest stones again and again, in a pile that changes after every round.

2. Hint 2: Ask a smaller question

   Take the stones 2, 7, 4, 1, 8, 1. Smash 8 and 7, which leaves a new stone of weight 1. Where does that stone stand among the others? To find the next two heaviest, do you need the whole pile in order, or only to know which stones are heaviest right now?

3. Hint 3: Name the pattern

   You need the heaviest stone again and again while the pile keeps changing. That is the signal. A max-heap always has its largest value on top, and adding a value or removing the largest takes about log n steps. Every stone waits in it and the heaviest is served first: this is a heap as a priority queue. Read the lesson on heap as a priority queue.

4. Hint 4: State the invariant

   At the start of every round the heap holds exactly the stones that are left, with the heaviest on top. So the first stone you take out is the heaviest, and the second is the heaviest of the rest.

5. Hint 5: Sketch the pseudocode

   Push every stone into a max-heap. While the heap holds at least two stones: pop the largest twice, calling them heavy and light. If heavy is not equal to light, push heavy minus light. When the loop ends, return the one stone left in the heap, or 0 if the heap is empty.

Solution

From here on the page gives the method away. If you are still solving, go back to the hints.

Approach

Put the stones in a max-heap. Take out the two heaviest and, if their weights differ, put the difference back. Repeat until at most one stone is left.

The slow way

Sort the stones, take the two heaviest off the end, add the difference if there is one, and sort again before the next round. It is correct, but there can be almost as many rounds as stones, and every round pays for a full sort of the pile, about n log n steps each.

The insight

A round only ever needs the two heaviest stones. The order of all the others does not matter yet. A max-heap is a container that keeps exactly that much order: its largest value is always on top, at the root, and adding a stone or removing the largest takes about log n steps instead of a full sort.

Step by step

1. Put every stone in a max-heap, which hands out the heaviest first. Every round asks for the heaviest stones that are left. A max-heap answers that in about log n steps, however the pile has changed since the last round.
2. Smash stones as long as at least two are left. A smash needs two stones. With one stone or none left there is nothing more to do, and the loop stops.
3. Take the two heaviest stones out of the heap. The heap hands out its largest value first, so the first stone taken is at least as heavy as the second. Their difference can then never be negative.
4. If their weights differ, put the difference back as a stone. Equal stones destroy each other and nothing is left of them. Otherwise the heavier one survives with the difference as its weight, and it has to go back into the heap to compete with the rest.
5. Return the last stone's weight, or 0 when no stone is left. Every round removes at least one stone, so the loop ends with one stone or none. One stone is the answer. None means no stone survived, and the statement asks for 0 then.

The same numbered steps are the comments in each solution.

function solve(stones) {
  // JavaScript has no built-in heap, so these two helpers keep an array in
  // max-heap order by hand: no value is larger than its parent, so the largest
  // is at index 0. The parent of index i is at (i - 1) >>> 1, which is
  // (i - 1) / 2 rounded down, and its two children are at 2 * i + 1 and
  // 2 * i + 2.
  function heapPush(heap, x) {
    // The new value starts at the end and swaps with its parent while the
    // parent is smaller.
    heap.push(x);
    let i = heap.length - 1;
    while (i) {
      const p = (i - 1) >>> 1;
      if (heap[p] >= heap[i]) break;
      [heap[p], heap[i]] = [heap[i], heap[p]];
      i = p;
    }
  }
  function heapPop(heap) {
    // The largest value leaves. The last value takes its place at index 0,
    // then swaps with its larger child while that child is larger.
    const largest = heap[0],
      last = heap.pop();
    if (!heap.length) return largest;
    heap[0] = last;
    let i = 0;
    while (true) {
      let child = 2 * i + 1;
      if (child >= heap.length) break;
      if (child + 1 < heap.length && heap[child + 1] > heap[child]) child++;
      if (heap[i] >= heap[child]) break;
      [heap[i], heap[child]] = [heap[child], heap[i]];
      i = child;
    }
    return largest;
  }
  // 1. Put every stone in a max-heap, which hands out the heaviest first
  const heap = [];
  for (const stone of stones) heapPush(heap, stone);
  // 2. Smash stones as long as at least two are left
  while (heap.length > 1) {
    // 3. Take the two heaviest stones out of the heap
    const heavy = heapPop(heap),
      light = heapPop(heap);
    // 4. If their weights differ, put the difference back as a stone
    if (heavy !== light) heapPush(heap, heavy - light);
  }
  // 5. Return the last stone's weight, or 0 when no stone is left
  return heap.length ? heap[0] : 0;
}

function solve(stones: number[]): number {
  // TypeScript has no built-in heap, so these two helpers keep an array in
  // max-heap order by hand: no value is larger than its parent, so the largest
  // is at index 0. The parent of index i is at (i - 1) >>> 1, which is
  // (i - 1) / 2 rounded down, and its two children are at 2 * i + 1 and
  // 2 * i + 2.
  function heapPush(heap: number[], x: number): void {
    // The new value starts at the end and swaps with its parent while the
    // parent is smaller.
    heap.push(x);
    let i = heap.length - 1;
    while (i) {
      const p = (i - 1) >>> 1;
      if (heap[p] >= heap[i]) break;
      [heap[p], heap[i]] = [heap[i], heap[p]];
      i = p;
    }
  }
  function heapPop(heap: number[]): number {
    // The largest value leaves. The last value takes its place at index 0,
    // then swaps with its larger child while that child is larger.
    const largest = heap[0],
      last = heap.pop()!;
    if (!heap.length) return largest;
    heap[0] = last;
    let i = 0;
    while (true) {
      let child = 2 * i + 1;
      if (child >= heap.length) break;
      if (child + 1 < heap.length && heap[child + 1] > heap[child]) child++;
      if (heap[i] >= heap[child]) break;
      [heap[i], heap[child]] = [heap[child], heap[i]];
      i = child;
    }
    return largest;
  }
  // 1. Put every stone in a max-heap, which hands out the heaviest first
  const heap: number[] = [];
  for (const stone of stones) heapPush(heap, stone);
  // 2. Smash stones as long as at least two are left
  while (heap.length > 1) {
    // 3. Take the two heaviest stones out of the heap
    const heavy = heapPop(heap),
      light = heapPop(heap);
    // 4. If their weights differ, put the difference back as a stone
    if (heavy !== light) heapPush(heap, heavy - light);
  }
  // 5. Return the last stone's weight, or 0 when no stone is left
  return heap.length ? heap[0] : 0;
}

import heapq


def solve(stones):
    # 1. Put every stone in a max-heap, which hands out the heaviest first
    # heapq always hands out the smallest number, so each weight is stored
    # negated: the smallest stored number then belongs to the heaviest stone.
    heap = []
    for stone in stones:
        heapq.heappush(heap, -stone)
    # 2. Smash stones as long as at least two are left
    while len(heap) > 1:
        # 3. Take the two heaviest stones out of the heap
        heavy = -heapq.heappop(heap)
        light = -heapq.heappop(heap)
        # 4. If their weights differ, put the difference back as a stone
        if heavy != light:
            heapq.heappush(heap, -(heavy - light))
    # 5. Return the last stone's weight, or 0 when no stone is left
    return -heap[0] if heap else 0

Example 1, traced

Input
stones = [2,7,4,1,8,1]

Output
1

Why it works

At the start of each round the heap holds exactly the stones that are left. Taking the largest twice gives the two heaviest, which is what the rules ask for, and the difference, if there is one, goes back before the next round. So the program plays the game move for move, and what is left at the end is the game's result.

Time and space

Time O(n log n), extra space O(n). A heap of n values has about log n levels, and adding or removing a value moves along them. There are n stones to put in and at most n rounds of a few such operations each: O(n log n). The heap holds at most n stones: O(n) space.

n and m are input sizes, h is a tree's height, V and E are a graph's vertices and edges. The bounds describe this reference solution.

Watch out

• Check that a stone is left before you read the top of the heap. When the answer is 0, as for [3, 3] or the empty list, the heap is empty: Python raises an error and JavaScript gives undefined.
• Put the new stone back with the heap's own push. Appended to the end of the list instead, it skips the reordering, the heap later hands out a stone that is not the heaviest, and example 1 returns -1.
• In Python, heapq hands out the smallest value, not the largest. Without the minus signs the loop takes the two lightest stones each round, and example 1 returns -17 instead of 1.

Pattern

Heap as a priority queue. Builds on Kth Largest Element In a Stream.

Keep going