Start learning free

LESSON 318 OF 414 · HEAP

Last Stone Weight in Rust: Simulating Smashes with a BinaryHeap

Practice O(n log n) Read the full explanation

We have a collection of stones, where each stone has a positive integer weight.

Each turn, we choose the two heaviest stones and smash them together.  Suppose the stones have weights x and y with x <= y.  The result of this smash is:

  • If x == y, both stones are totally destroyed;
  • If x != y, the stone of weight x is totally destroyed, and the stone of weight y has the new weight y-x.

At the end, there is at most 1 stone left.  Return the weight of this stone (or 0 if there are no stones left.)

 

Example:

Input: [2,7,4,1,8,1]
Output: 1
Explanation: 
We combine 7 and 8 to get 1 so the array converts to [2,4,1,1,1] then,
we combine 2 and 4 to get 2 so the array converts to [2,1,1,1] then,
we combine 2 and 1 to get 1 so the array converts to [1,1,1] then,
we combine 1 and 1 to get 0 so the array converts to [1] then that's the value of the last stone.
STEP 1 OF 6
    Ctrl + Enter to check

    RUST · HEAP

    Last Stone Weight in Rust, explained

    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

    1. Build the heap. Declare a mutable BinaryHeap<i32> named heap by collecting stones.into_iter() into it. The iterator hands over the stones by value, and collect builds the heap from them.
    2. Loop while at least two stones remain. The condition heap.len() >= 2 means a smash is possible. When the loop ends, the heap has zero or one stone.
    3. Take the heaviest stone into x with heap.pop().unwrap(). pop returns an Option, and the loop condition guarantees it is not empty, so unwrap can't fail.
    4. Take the second heaviest into y the same way. Since x came out first, x >= y.
    5. Push the difference when the stones differ. If x != y, put x - y back. Equal stones destroy each other, so nothing goes back.
    6. 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 on heap says which one. For a BinaryHeap, this builds the heap in one pass, which takes O(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() against unwrap_or(0). unwrap() panics on None. It is safe inside the loop because the length check came first. At the end the heap might be empty, so unwrap_or(0) supplies the default instead of panicking.
    • No overflow in x - y. Both are positive and x >= y, so the result is between 0 and x. A difference of 0 can't happen, because equal stones are handled by the if and 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 second unwrap would panic when only one stone is left.
    • Pushing 0. Skipping the x != y test 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] gives 1.
    • stones = [10, 4] gives 6.
    • stones = [3, 3] gives 0.
    • stones = [1, 1, 1] gives 1.
    • stones = [5] gives 5.
    • stones = [] gives 0.
    Keep going with RustAbout 340 step-by-step lessons, from your first line of code to coding interviews. Free to start.
    Start learning free