Smash the two heaviest stones together until at most one is left. A max-heap hands you the two biggest every turn in log time.
▼The problem
LeetCode 1046 (Easy). You have a pile of stones with positive integer weights. Each turn, take the two heaviest, x ≤ y, and smash them together: if x = y both are destroyed, otherwise x is destroyed and y becomes y − x. When at most one stone is left, return its weight, or 0 if none are left.
Examples (LeetCode's): [2,7,4,1,8,1] → 1 (8 and 7 leave 1, 4 and 2 leave 2, 2 and 1 leave 1, 1 and 1 are both destroyed, leaving one stone of weight 1), and [1] → 1 (nothing to smash).
The solution
import heapq
def lastStoneWeight(stones):
heap = [-s for s in stones] # max-heap
heapq.heapify(heap)
while len(heap) > 1:
y = -heapq.heappop(heap) # heaviest
x = -heapq.heappop(heap) # next one
if y != x:
heapq.heappush(heap, -(y - x))
return -heap[0] if heap else 0Transcript
Last Stone Weight. You're given a pile of stones. Each turn, smash the two heaviest together. If they weigh the same, both are destroyed. If not, the lighter one is destroyed, and the heavier loses that much weight. Return the last stone's weight, or zero if none are left.
Take stones two, seven, four, one, eight, one. Eight and seven leave one. Four and two leave two. Two and one leave one. One and one destroy each other, so the last stone weighs one. A pile of just one stone has nothing to smash: the answer is one.
The naive way sorts the whole pile every turn: up to n sorts, each n log n. Even scanning for the top two costs order n squared.
The fix: a max heap, a pile where the heaviest stone is always on top. Popping the top, or pushing a stone, costs only log n, because a stone moves one level at a time.
Python's heapq is a min heap, so store negative weights. Heapify. While two or more stones remain, pop the two heaviest, and if they differ, push back the difference. Return the last stone, or zero.
Heapify the six stones, and eight rises to the top. Pop eight and seven, push one. Pop four and two, push two. Pop two and one, push one. Pop one and one: equal, so nothing goes back. One stone is left: one.
Each turn destroys at least one stone, so there are at most n turns of log n each. That's order n log n time, and order n space for the heap.
Pop the two heaviest, push back the difference. That's Last Stone Weight.