Problem overview
The function last_stone_weight receives a vector of positive integers, stones, the weights of some stones. Each turn takes the two heaviest stones and smashes them together. If they weigh the same, both disappear. If they differ, the lighter one disappears and the heavier one stays with the difference of the two weights. The turns repeat until at most one stone is left. The function returns the weight of that stone, or 0 if none is left.
For [2, 7, 4, 1, 8, 1] the process goes like this:
- 8 and 7 smash and leave 1, so the stones are 2, 4, 1, 1, 1;
- 4 and 2 leave 2, so the stones are 2, 1, 1, 1;
- 2 and 1 leave 1, so the stones are 1, 1, 1;
- 1 and 1 destroy each other, which leaves a single 1.
The answer is 1.
Approach: always take the largest two
Each turn needs the two largest weights, and the weight that remains goes back among the others, where it may become one of the largest again or not. Sorting the vector once isn't enough, because the new stone has to be inserted at the right place. Sorting again each turn works but costs O(n log n) per turn.
A max-heap fits the job. It gives the largest element and removes it in O(log n), and inserting a new element is also O(log n). Rust's BinaryHeap is a max-heap by default, so the heaviest stone is always the next one out, and no custom ordering is needed.
Every turn removes at least one stone from the heap, so there are fewer than n turns. Each turn does two pops and at most one push.
Step by step
- Build the heap. Declare a mutable
BinaryHeap<i32>namedheapby collectingstones.into_iter()into it. The iterator hands over the stones by value, andcollectbuilds the heap from them. - Loop while at least two stones remain. The condition
heap.len() >= 2means a smash is possible. When the loop ends, the heap has zero or one stone. - Take the heaviest stone into
xwithheap.pop().unwrap().popreturns anOption, and the loop condition guarantees it is not empty, sounwrapcan't fail. - Take the second heaviest into
ythe same way. Sincexcame out first,x >= y. - Push the difference when the stones differ. If
x != y, putx - yback. Equal stones destroy each other, so nothing goes back. - Return the last stone. After the loop,
heap.pop().unwrap_or(0), written as the final expression with no semicolon, returns the remaining stone, or 0 for an empty heap.
The solution in Rust
use std::collections::BinaryHeap;
impl Solution {
fn last_stone_weight(stones: Vec<i32>) -> i32 {
let mut heap: BinaryHeap<i32> = stones.into_iter().collect();
while heap.len() >= 2 {
let x = heap.pop().unwrap();
let y = heap.pop().unwrap();
if x != y {
heap.push(x - y);
}
}
heap.pop().unwrap_or(0)
}
}
The Rust details worth a closer look:
- Building a heap from an iterator.
collect()can build any type that knows how to be created from an iterator, and the type annotation onheapsays which one. For aBinaryHeap, this builds the heap in one pass, which takesO(n)time and is cheaper than pushing the stones one at a time. - Ownership.
stones.into_iter()consumes the vector, which is fine because the parameter is passed by value and nothing else uses it afterwards. pop().unwrap()againstunwrap_or(0).unwrap()panics onNone. It is safe inside the loop because the length check came first. At the end the heap might be empty, sounwrap_or(0)supplies the default instead of panicking.- No overflow in
x - y. Both are positive andx >= y, so the result is between 0 andx. A difference of 0 can't happen, because equal stones are handled by theifand never pushed. - Max-heap by default. For the opposite order, values would be wrapped in
std::cmp::Reverse. This lesson wants the largest first, so no wrapper is needed.
Complexity
Time: O(n log n). Building the heap takes O(n). Each turn removes at least one stone from the heap, because two come out and at most one goes back, so there are fewer than n turns. Each turn costs two pops and at most one push, each O(log n).
Space: O(n). The heap holds the stones. It never grows beyond n elements.
Edge cases and common mistakes
- No stones. An empty vector never enters the loop, and
unwrap_or(0)returns 0. - One stone.
[5]returns 5, since nothing smashes. - Two equal stones.
[3, 3]returns 0. The stones destroy each other. - Everything cancels.
[9, 3, 2, 10]returns 0: 10 and 9 leave 1, then 3 and 2 leave 1, then the two 1s destroy each other. - Odd count of equal stones.
[1, 1, 1]returns 1: one pair cancels and one stone stays. - Using
pop().unwrap()at the end. That panics for an empty heap, for example with[3, 3]. - Looping while
heap.len() > 0. The secondunwrapwould panic when only one stone is left. - Pushing 0. Skipping the
x != ytest puts a stone of weight 0 into the heap. The final weight is not affected, but the stone breaks the rule that equal stones both disappear, and it adds extra turns.
Testing
stones = [2, 7, 4, 1, 8, 1]gives1.stones = [10, 4]gives6.stones = [3, 3]gives0.stones = [1, 1, 1]gives1.stones = [5]gives5.stones = []gives0.
