parikshan

last-stone-weight

Difficulty: Easy  ·  Topic: Greedy  ·  Practice it on parikshan: /practice/last-stone-weight/start

The story

You build the matchmaking engine at a poker site like PokerStars. To run a "shootout" tournament you repeatedly pair the two highest chip stacks at a table; whoever has more wins, the loser is eliminated, and the winner's stack is reduced by the loser's stack (the chips taken). When two equal stacks meet, both bust (a split pot that consumes both stacks). The tournament runs until at most one player remains. The product team wants to know, before a tournament starts, what the final chip stack will look like, so they can size the prize pool against expected variance.

That is last-stone-weight. The simulation appears in dozens of guises: rocket-stage fuel pairing for spent boosters, antimatter collider experiments where two beams cancel by their mass difference, retail merger games where two companies of equal capitalisation cancel out. The defining shape is "repeatedly pick the two largest, replace by their difference (or destroy both if equal), until one or zero remain."

This is a wonderfully clean problem on which to internalise the right data structure. The greedy strategy is obvious, always pick the two largest. The whole craft is in choosing the structure that makes "two largest" cheap to query and cheap to update.

What's actually being asked

Given n stones with positive integer weights, repeatedly:

  1. Pick the two heaviest, y >= x.
  2. If y == x, both are destroyed.
  3. If y > x, the smaller is destroyed and the heavier becomes a new stone of weight y - x.

Stop when zero or one stone remains. Return the final weight, or 0 if empty.

The naive approach

Sort each round. Sort the array, pop the two largest, push the difference back, sort again. Each round is O(n log n) and there are up to n rounds (each round removes at least one stone), so the total is O(n^2 log n). For n = 10^5 that is 10^10 * log, utterly hopeless.

Linear scan each round. Scan to find the two largest in O(n), splice them out, push the new stone, scan again. O(n^2) total. For n = 10^5, that is 10^10, still hopeless.

The right answer is to use a structure that gives you the largest in O(log n) per query and per update. That structure is a heap.

The key insight

The pattern "give me the largest, modify it, put it back, repeat" is the signature use case of a priority queue. Each pop, push, and peek is O(log n). With at most n rounds and two pops + one optional push per round, the total work is O(n log n). That is the same complexity as the initial sort and far better than O(n^2 log n).

This is the lesson last-stone-weight teaches: the algorithm is trivial; the data structure is everything. Recognising "I need the max under updates" is the cue to reach for a heap. Once you have the cue in your fingers, problems like k-closest-points-to-origin, merge-k-sorted-lists, task-scheduler, and meeting-rooms-ii all become five-minute exercises.

The greedy correctness argument

Why is "always pick the two largest" correct? Intuition first: the two largest cancel as much weight as possible in a single round (because either both are destroyed or the smaller is, and the residue y - x is the smallest the larger can become in one move). Smaller stones cannot affect the outcome of the two largest because their interaction does not change weights above them in the rank order. So the final residue depends only on the cancellation pattern at the top of the heap.

A formal proof is by induction on n. Base: n <= 1 is trivial. Inductive step: assume the greedy answer is optimal for all sizes below n. For size n, the first move pairs the top two; that choice removes one or two stones from the heap. The remaining problem has size at most n - 1, so by hypothesis greedy completes it correctly. The first move was the best because any other pair would leave a larger stone in play, which dominates the outcome (a strictly larger top stone produces a strictly larger or equal final residue when the smaller-paired alternative ran). The full statement and case analysis is more delicate, but the intuition is right.

For interview purposes, the greedy choice is so natural that the focus is correctly on the heap. State the choice, name the proof family ("exchange-style on the largest pair"), and move on.

Python's heapq trick: negate for max-heap

heapq in the standard library is a min-heap. To get a max-heap, push the negatives of weights. Pop returns the smallest negative, which is the largest original. Negate again on the way out. This is idiomatic and zero-cost.

Step-by-step approach

  1. Build a list of negated weights: h = [-w for w in stones].
  2. heapq.heapify(h). This is O(n), faster than n pushes.
  3. While len(h) >= 2:
    • y = -heapq.heappop(h); x = -heapq.heappop(h). Because we negated, the two smallest negatives are the two largest originals; y >= x by heap invariant.
    • If y > x: heapq.heappush(h, -(y - x)).
    • If y == x: push nothing, both are destroyed.
  4. If h is empty, print 0. Else print -h[0].

Worked example

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

Negated, heapified: [-8, -7, -4, -1, -2, -1] (heap order, not sorted).

  • Pop two: y=8, x=7. Diff 1, push -1. Heap now represents [4, 2, 1, 1, 1].
  • Pop two: y=4, x=2. Diff 2, push -2. Heap represents [2, 1, 1, 1].
  • Pop two: y=2, x=1. Diff 1, push -1. Heap represents [1, 1, 1].
  • Pop two: y=1, x=1. Equal, push nothing. Heap represents [1].
  • Loop ends, len(h) = 1. Print -h[0] = 1.

Output: 1. Matches the spec.

Complexity

  • Time: O(n log n). Heapify is O(n); each of at most n rounds does two O(log n) pops and at most one O(log n) push.
  • Space: O(n) for the heap.

You cannot do better in the comparison model: even reading the input is O(n), and any algorithm that interacts with all stones must inspect them at least once.

Common pitfalls

  • Forgetting to negate: pushing positive weights into heapq gives a min-heap, and you would always pull the two smallest. Wrong answer, easy to miss because the algorithm "looks right".
  • Pushing the difference when it is zero: pushing -0 is 0, a valid heap element, which would create a phantom stone of weight 0. The spec says equal stones are destroyed; check y > x before pushing.
  • Off-by-one on the loop condition: while len(h) >= 2 is correct. while len(h) > 1 is equivalent. while h is wrong, you stop when only one element remains, not zero.
  • Using sorted() instead of heapify: sorted is O(n log n) and gives you a sorted list, but each subsequent insertion is O(n). The heap interface is what you want.

Where this shows up in the enterprise

  • PokerStars / WSOP.com: shootout-style tournament simulators, where the top two stacks at each table collide until one remains.
  • CERN / Fermilab: pair-cancellation experiments where two particles annihilate by mass difference, simulation tools for collider event triage.
  • Lockheed Martin / SpaceX: spent-stage propellant balancing simulators; pair the heaviest fuel remainders for vent operations.
  • Goldman Sachs / Renaissance Technologies: position-cancellation simulators in equity arbitrage where two large opposing positions cancel by net delta.
  • AWS Auto Scaling: pairing the two heaviest pending jobs to a worker is a heap operation on job size; the residue queues again if work remains.
  • Snowflake / BigQuery: query optimisers' equi-join planners use heaps for selecting the two largest unjoined relations.
  • Datadog: alert deduplication pairs the two highest-priority signals from the same source, keeping the residue priority.

The universal pattern is "merge or cancel by priority". Whenever you see "always pick the largest" plus an update, the answer is a heap. That is the take-away worth more than the problem itself.

Variants you'll see elsewhere

  • kth-largest-element: maintain a heap of size k; same trick.
  • k-closest-points-to-origin: heap of size k keyed by distance.
  • merge-k-sorted-lists: heap of k list heads, repeatedly pop and advance.
  • task-scheduler: heap of task frequencies; pop the most frequent, cooldown, push back.
  • meeting-rooms-ii: heap of room end-times; check if the earliest-ending room is free for each new meeting.
  • find-median-from-data-stream: two heaps (max-heap of lower half, min-heap of upper half), keeping sizes balanced.

Six problems, one data structure. Learn the heap once on last-stone-weight and you have the key to all of them.

How parikshan turns this into a learning loop

The dynamic private tests cover n = 1, all-equal even count (cancels to 0), all-equal odd count (one stone of original weight left), and stress inputs at n = 10^5 where the O(n^2) solutions time out. The hint ladder explicitly names the heap at hint 1, the negation trick at hint 3, and the full skeleton at hint 4, a generous progression because the focus is on choosing the structure, not deriving the algorithm. Practice mode is for asking the AI assistant to Explain Python's heapq interface and the negation idiom. Exam mode is where you write the seven lines from memory at speed.

In the AI-integrated workspace

Agents handle this problem competently. The interesting failures are at the level of structure choice, not implementation:

  1. Agent reaches for sorting, not a heap. The generated code looks fine on small inputs and TLEs on stress. The fix is recognising the pattern; the code is mechanical once you have it.
  2. Agent forgets the negation when porting from a language with native max-heaps (C++ priority_queue, Java PriorityQueue with a reverse comparator). The Python translation needs negation; if you do not double-check, you ship a min-heap that gives the wrong answer.
  3. Agent writes a custom MaxHeap class to avoid the negation trick. Three times the code, no faster, and a fresh source of bugs. The negation idiom is the Pythonic answer; the custom class is over-engineering.

The senior-engineer move on heap problems is to name the structure explicitly in the prompt: "use heapq with negated weights". That single sentence prevents all three failure modes. Once the structure is fixed, the rest is six lines of mechanical code.