Start learning free

LESSON 318 OF 415 · HEAP

Last Stone Weight in Swift: Simulating Smashes with a Max Heap

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 8

Guided

Ctrl + Enter to check

SWIFT · HEAP

Last Stone Weight in Swift, explained

Problem overview

The method lastStoneWeight receives an array of positive integers, stones, the weights of some stones. In every turn the two heaviest stones are smashed together. If they have the same weight, both disappear. Otherwise the lighter one disappears and the heavier one stays with a new weight, the difference between the two. The process repeats until at most one stone is left. The method returns that stone's weight, or 0 if no stone is left.

Take [2, 7, 4, 1, 8, 1]. The stones 8 and 7 give a stone of weight 1, leaving [2, 4, 1, 1, 1]. Then 4 and 2 give 2, then 2 and 1 give 1, then 1 and 1 destroy each other, and one stone of weight 1 remains. The answer is 1.

Approach: always ask for the heaviest

Every turn needs the two largest values of a collection that keeps changing. Sorting the array again after every smash would work, but it is wasteful. A max heap gives the largest element in O(1), removes it in O(log n) and accepts a new element in O(log n).

So the plan is to load all the stones into a max heap and repeat this while at least two stones remain: take out the two heaviest, and if their weights differ, put the difference back. Each turn removes at least one stone from the heap, so the loop ends after at most n - 1 turns.

Swift's standard library has no heap, so the lesson provides a MaxHeap struct, ready to use. The lesson uses its methods heapify, extract, insert and count.

Step by step

Create the heap. var heap = MaxHeap(). It must be a var because its methods modify the struct.
Load the stones. heap.heapify(stones) puts every stone into the heap, so the heaviest is on top.
Loop while a smash is possible. The condition is heap.count >= 2. When the loop ends, at most one stone is left.
Take the heaviest. let x = heap.extract() removes the top and returns it.
Take the second heaviest. let y = heap.extract(). Since x came out first, x >= y.
Put back the remainder. If x != y, call heap.insert(x - y). If the weights are equal, both stones are gone and nothing goes back.
Handle no stones left. After the loop, if heap.count == 0, return 0.
Return the survivor. Otherwise exactly one stone remains, and heap.extract() returns it.

The solution in Swift

class Solution {
    func lastStoneWeight(_ stones: [Int]) -> Int {
        var heap = MaxHeap()
        heap.heapify(stones)
        while heap.count >= 2 {
            let x = heap.extract()
            let y = heap.extract()
            if x != y {
                heap.insert(x - y)
            }
        }
        if heap.count == 0 {
            return 0
        }
        return heap.extract()
    }
}

The Swift details worth a closer look:

var for a struct. MaxHeap is a struct, which is a value type. Methods that change it, such as extract and insert, are mutating, and Swift allows calling them only on a var. Declaring heap with let would be a compile error.
let x and let y. Both values are assigned once and never change, so let is the right choice. Swift code uses let by default and var only when a value must change.
count is a property. Write heap.count, not heap.count().
Extracting from an empty heap. Taking a value out of an empty heap is not meaningful and can stop the program, so the code never does it. That is why the loop checks count >= 2 before taking two values, and why the final extract() comes after the count == 0 check.
No parameter labels. The method is declared with _ in front of stones, so the call is lastStoneWeight([2, 7, 4, 1, 8, 1]) without a label.

Complexity

Time: O(n log n). The loop runs at most n - 1 times, and each turn does two extractions and at most one insertion, each costing O(log n). Loading the stones costs no more than inserting them one by one, which is also O(n log n).

Space: O(n). The heap holds the stones, and it never holds more than the original n.

Edge cases and common mistakes

Empty array. The loop doesn't run, heap.count is 0, and the method returns 0.
One stone. [5] returns 5, because no smash happens.
Two equal stones. [3, 3] returns 0.
Several equal stones. [1, 1, 1] returns 1: two of them destroy each other and one is left.
A chain of smashes. [9, 3, 2, 10] smashes 10 and 9 into a 1, leaving 3, 2 and 1. Then 3 and 2 give another 1, and the two 1s destroy each other. The result is 0. Note that the new stone goes back into the heap and can be picked again at once, or left waiting, depending on its weight.
Inserting y - x. The first value taken out is the larger one, so the remainder is x - y. The reverse order would put a negative weight in the heap.
Inserting 0 on equal weights. A stone of weight 0 would stay in the heap and be counted as a stone, so only insert when x != y.

Testing

[2, 7, 4, 1, 8, 1] gives 1.
[] gives 0.
[5] gives 5.
[3, 3] gives 0.
[9, 3, 2, 10] gives 0.

Keep going with SwiftAbout 340 step-by-step lessons, from your first line of code to coding interviews. Free to start.

Start learning free