Start learning free

LESSON 287 OF 338 · HEAP

Last Stone Weight in Go: O(n log n) 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 7
    Ctrl + Enter to check

    GO · HEAP

    Last Stone Weight in Go, explained

    Problem overview

    The function lastStoneWeight receives a slice stones of positive integer weights. Each turn, the two heaviest stones are smashed together. If they weigh the same, both disappear. If they differ, the lighter one disappears and the heavier one stays with the difference as its new weight. The game ends when at most one stone is left, and the function returns that stone's weight, or 0 if no stone is left.

    In the example, stones = []int{2, 7, 4, 1, 8, 1}:

    • 8 and 7 are smashed and leave a stone of weight 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 two largest

    The process says exactly what to do each turn: find the two heaviest stones, remove them, and put back their difference if it isn't zero. The question is how to find the heaviest stones quickly.

    Sorting after every smash costs O(n log n) per turn, and the game can last n - 1 turns. Scanning for the two largest each time is O(n) per turn. A max heap is better suited: it gives the largest element in O(1) and removes it, or inserts a new element, in O(log n).

    The plan:

    1. Turn the slice into a max heap.
    2. While at least two stones remain, pop the two largest, x then y. Since x came out first, x >= y.
    3. If x != y, push x - y back. If they are equal, both stones are gone and nothing is pushed.
    4. Each turn removes at least one stone, so the loop ends. At most one stone remains: return it, or 0 if the heap is empty.

    Step by step

    1. Convert stones to a maxHeap named stoneHeap and call heap.Init(&stoneHeap). The maxHeap type is provided by the lesson: a slice of ints that implements heap.Interface with Less reversed, so the largest value is on top. heap.Init rearranges the whole slice into heap order in O(n), which is cheaper than pushing the stones one at a time.
    2. Loop while there are at least two stones: for stoneHeap.Len() >= 2. Go has no while, and a for with only a condition does the job.
    3. Pop the heaviest stone into x. heap.Pop returns an any, so the type assertion .(int) turns it back into an int.
    4. Pop the next heaviest into y the same way.
    5. Push the difference if there is one. If x != y, call heap.Push(&stoneHeap, x - y). Use the package function, not the method on the type, so the heap order is restored.
    6. Handle the empty case. After the loop, if stoneHeap.Len() == 0, return 0.
    7. Return the last stone. Otherwise exactly one stone is left, and it sits at index 0, the root of the heap.

    The solution in Go

    The solution uses the provided maxHeap type together with the container/heap package. A typical definition of that type is:

    type maxHeap []int

func (h maxHeap) Len() int           { return len(h) }
func (h maxHeap) Less(i, j int) bool { return h[i] > h[j] }
func (h maxHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h *maxHeap) Push(x any)        { *h = append(*h, x.(int)) }
func (h *maxHeap) Pop() any {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[:n-1]
    return x
}

    And the function:

    func lastStoneWeight(stones []int) int {
    stoneHeap := maxHeap(stones)
    heap.Init(&stoneHeap)
    for stoneHeap.Len() >= 2 {
        x := heap.Pop(&stoneHeap).(int)
        y := heap.Pop(&stoneHeap).(int)
        if x != y {
            heap.Push(&stoneHeap, x - y)
        }
    }
    if stoneHeap.Len() == 0 {
        return 0
    }
    return stoneHeap[0]
}

    The Go details worth knowing:

    • Type conversion shares memory. maxHeap(stones) converts a []int to the named type with the same underlying slice. Nothing is copied, so heap.Init and the later pops rearrange and overwrite the caller's slice. After the example call, the caller's stones reads [1 1 1 2 4 8]. Copy the slice first if the caller needs the original order.
    • Why &stoneHeap. Push and Pop change the slice's length, so they have pointer receivers, and heap.Init, heap.Pop and heap.Push all take the pointer. stoneHeap.Len() and stoneHeap[0] work on the variable directly.
    • Package functions, not methods. The type's own Push and Pop only append or cut the last element. The functions heap.Push and heap.Pop call them and restore the heap order.
    • Type assertion. heap.Pop(...).(int) panics if the stored value isn't an int. Here it always is.
    • Index 0 is the root. In a heap stored in a slice, the top is at the front. When only one element is left, it is the answer.

    Complexity

    Let n be the number of stones.

    Time: O(n log n). Building the heap with heap.Init takes O(n). Every turn removes at least one stone, so there are at most n - 1 turns. Each turn has two pops and at most one push, each O(log n).

    Space: O(1) extra. The heap is built inside the input slice. The loop pushes at most one stone after popping two, so the slice never needs to grow.

    Edge cases and common mistakes

    • Empty slice. The loop doesn't run, the length is 0, and the function returns 0.
    • One stone. stones = []int{9} returns 9, because there is nothing to smash.
    • Two equal stones. stones = []int{3, 3} returns 0.
    • Several rounds. stones = []int{10, 4, 2, 10} returns 2: the two 10s cancel, then 4 and 2 leave 2.
    • Pushing zero. Pushing x - y even when it is zero would leave a phantom stone of weight 0 in the heap, and the length checks would be wrong. The if x != y guard prevents it.
    • Using the heap as a min heap. If Less used <, the lightest stones would be smashed first.
    • Popping with the method. stoneHeap.Pop() removes the last slice element, not the largest stone.
    • Forgetting heap.Init. The slice is in arbitrary order, so the first pops would not return the heaviest stones.

    Testing

    • stones = []int{2, 7, 4, 1, 8, 1} gives 1.
    • stones = []int{3, 3} gives 0.
    • stones = []int{9} gives 9.
    • stones = []int{10, 4, 2, 10} gives 2.
    • stones = []int{} gives 0.
    Keep going with GoAbout 340 step-by-step lessons, from your first line of code to coding interviews. Free to start.
    Start learning free