DSAMotivator

Last Stone Weight

Easy·⛰️ Heaps / Priority Queue·Open on LeetCode ↗

Repeatedly take the two heaviest stones and smash them together — equal stones destroy each other, unequal ones leave a remnant — until at most one stone remains.

stones = [5, 3, 9, 2, 4]step 1/16
5
heap array
0
5

Add stone 5 — heavier than everything here, so it rises to the top.

1import heapq
2 
3def last_stone_weight(stones):
4    heap = [-s for s in stones]
5    heapq.heapify(heap)
6    while len(heap) > 1:
7        a = -heapq.heappop(heap)
8        b = -heapq.heappop(heap)
9        if a != b:
10            heapq.heappush(heap, -(a - b))
11    return -heap[0] if heap else 0
Read the 16 steps as text
  1. 1Add stone 5 — heavier than everything here, so it rises to the top.
  2. 2Add stone 3. It bubbles up until its parent outweighs it.
  3. 3Add stone 9 — heavier than everything here, so it rises to the top.
  4. 4Add stone 2. It bubbles up until its parent outweighs it.
  5. 5Add stone 4. It bubbles up until its parent outweighs it.
  6. 6Heap built. Every parent outweighs its children — so the heaviest stone is the root, always.
  7. 7Take the root, 9. The last stone moves up and sinks back to its place.
  8. 8Take the new root, 5. Two heaviest in hand.
  9. 99 − 5 = 4. The remnant goes back in and rises to where it belongs.
  10. 10Take the root, 4. The last stone moves up and sinks back to its place.
  11. 11Take the new root, 4. Two heaviest in hand.
  12. 124 and 4 are equal — both destroyed, nothing goes back.
  13. 13Take the root, 3. The last stone moves up and sinks back to its place.
  14. 14Take the new root, 2. Both stones are out, so the heap is momentarily empty.
  15. 153 − 2 = 1. The remnant goes back in and rises to where it belongs.
  16. 16One stone left. The answer is 1.

The idea

The problem hands you its own algorithm: take the two heaviest, smash, repeat. The only real question is how to keep finding 'the two heaviest' when the set keeps changing underneath you.

Re-sorting after every smash works and is O(n² log n). But you never need the whole order — you need the maximum, twice, and then you need to insert one new element. That exact trio of operations is what a heap is for: O(1) to peek at the max, O(log n) to remove it, O(log n) to insert.

A max-heap is a binary tree with one rule — every parent is at least as heavy as its children — and it's stored in a flat array, where node i has children at 2i+1 and 2i+2. That rule is deliberately weaker than sorting. It says nothing about siblings, which is precisely why maintaining it costs log n instead of n log n. You're paying only for the ordering you actually use.

Python's heapq is a min-heap and has no reverse flag, so the standard move is negating every value on the way in and on the way out. It's not elegant, but interviewers expect it and its absence reads as unfamiliarity with the library.

The approach

  1. 1Put every stone into a max-heap.
  2. 2While more than one stone remains, remove the two largest.
  3. 3If they differ, push their difference back in; if they're equal, push nothing — both are gone.
  4. 4When the loop ends, the heap holds either the last stone or nothing at all.

Complexity

Time
O(n log n)
Space
O(n)

Building the heap is O(n) with heapify. Each smash removes at least one stone and costs O(log n), so there are at most n rounds. The space is the heap itself.

Code

import heapq

def last_stone_weight(stones):
    heap = [-s for s in stones]
    heapq.heapify(heap)
    while len(heap) > 1:
        a = -heapq.heappop(heap)
        b = -heapq.heappop(heap)
        if a != b:
            heapq.heappush(heap, -(a - b))
    return -heap[0] if heap else 0

What goes wrong

  • •Returning `heap[0]` without handling the empty case. If the stones annihilate perfectly the heap ends up empty and the answer is 0, not a crash.
  • •Forgetting to negate on the way *out* in Python. The values come back negative and the function returns something like -1, which passes a surprising number of hand-tests because small cases often end at 0.
  • •Pushing the difference when it's zero. It isn't wrong for the final answer, but it puts a phantom 0-weight stone in the heap that can be selected as one of the 'two heaviest' later — and on the equal-stones path it changes the loop's termination.
  • •Reaching for `sorted()` inside the loop. It gives the right answer and is the honest fallback if you blank on the heap API, but say out loud that it's O(n² log n) — being asked to improve it is worse than pre-empting it.