Last Stone Weight (LeetCode 1046)

The problem

You get a list of stone weights. Each turn, take the two heaviest stones, a ≥ b, and smash them together:

• if a == b, both are destroyed;
• otherwise b is destroyed and a becomes a - b.

Repeat until at most one stone is left. Return its weight, or 0 if none remain.

Input:  [2, 7, 4, 1, 8, 1]
Output: 1

8 and 7 leave 1; then 4 and 2 leave 2; then 2 and 1 leave 1; then 1 and 1 cancel. One stone of weight 1 remains.

How to think about it

Start with brute force. Each turn, sort the list (or scan it twice) to find the two largest, smash them, and put the result back. Sorting every round is O(n log n) per turn and there can be n turns, so O(n² log n). Even with two linear scans it’s O(n²).

The key insight: you only ever need the current maximum. You keep asking “what’s the biggest right now?” while the collection changes. That’s exactly what a max-heap is for: pop the largest in O(log n), push a new value in O(log n). Each turn becomes two pops and at most one push.

Python only has a min-heap. heapq always pops the smallest item. The standard trick is to store negated weights: the smallest negative is the largest weight. Negate on the way in and on the way out.

Which pattern is this? Heaps. The tell is “repeatedly take the largest (or smallest) from a collection that keeps changing”.

Approach

1. Negate every weight and heapify the list: O(n).
2. While there are at least two stones, pop the two largest, a then b.
3. If they differ, push back a - b (negated).
4. When the loop ends, return the remaining stone’s weight, or 0 if the heap is empty.

Solution

import heapq
def last_stone_weight(stones):
    h = [-s for s in stones]; heapq.heapify(h)               # negate: min-heap acts as max-heap
    while len(h) > 1:
        a = -heapq.heappop(h); b = -heapq.heappop(h)         # two heaviest, a >= b
        if a != b: heapq.heappush(h, -(a - b))               # leftover goes back in
    return -h[0] if h else 0
# last_stone_weight([2,7,4,1,8,1]) -> 1

Dry run

On [2, 7, 4, 1, 8, 1] (weights shown positive, sorted for readability):

One stone is left, so the answer is 1.

Complexity

• Time: O(n log n). Heapify is O(n). Every turn removes at least one stone, so there are at most n - 1 turns, each doing O(log n) heap work.
• Space: O(n). The heap holds a copy of the stones.

Edge cases and pitfalls

• One stone is returned as-is: [5] gives 5.
• All destroyed, as with [2, 2], gives 0. That’s why the return checks if h.
• Forgetting to negate back. -heapq.heappop(h) is easy to write as heapq.heappop(h), which hands you a negative weight and silently breaks the subtraction.
• Order of the two pops. The first pop is the larger, so a - b is always positive. Swapping them pushes a wrong sign.

Related problems

• Kth Largest Element in an Array (LeetCode 215): the next step with heaps, keeping only k items.
• Task Scheduler (LeetCode 621): another “most frequent first” simulation, here solved with a formula.
• Reorganize String (LeetCode 767): a max-heap simulation with a cooldown twist.