Smash Heaviest Stones
Problem
A pile of stones is described by their weights. Repeatedly take the two heaviest stones and smash them together: equal stones destroy each other completely, while unequal ones leave a single stone weighing the difference, which goes back into the pile. Return the weight of the last stone left, or 0 when the pile empties.
Examples
Input: stones = [2, 7, 4, 1, 8, 1]
Output: 1
Why: 7 and 8 leave 1, then 2 and 4 leave 2, then 1 and 2 leave 1, then 1 and 1 vanishInput: stones = [1, 1]
Output: 0
Why: two equal stones destroy each other and nothing remainsInput: stones = [3]
Output: 3
Why: edge case, a lone stone has nothing to smash againstHints
0 / 3
Re-sorting the whole pile after every smash is correct but does far more work than the question needs, since only two stones change each round.
The only thing ever asked of the pile is give me the current heaviest, and the only thing ever added back is a single new stone. That pair of operations names the structure.
Load every weight into a structure that always hands back the largest item and accepts new items cheaply. Each round, take two items, and put their difference back only when it is not zero. Stop when fewer than two items remain.
Solution
Only two questions are ever asked of the pile, give me the heaviest and take this new stone, which is exactly what a heap answers in logarithmic time. Building it once and then smashing repeatedly avoids re-sorting after every round. A difference of zero is simply not pushed back, so the pile shrinks by two in that case and by one otherwise, and the loop ends with either one stone or none. Time is O(n log n), and space is O(n) for the heap.
import heapq
def last_stone_weight(stones):
heap = [-w for w in stones] # negated, because the library heap is a min-heap
heapq.heapify(heap)
while len(heap) > 1:
heaviest = -heapq.heappop(heap)
second = -heapq.heappop(heap)
if heaviest != second: # equal stones vanish, so nothing goes back
heapq.heappush(heap, -(heaviest - second))
return -heap[0] if heap else 0
print(last_stone_weight([2, 7, 4, 1, 8, 1])) # -> 1
print(last_stone_weight([1, 1])) # -> 0
print(last_stone_weight([3])) # -> 3Stuck on the idea rather than the code? Heap Basics covers it.