Start learning free

LESSON 293 OF 348 · HEAP

Last Stone Weight in Kotlin: Simulating 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
    Ctrl + Enter to check

    KOTLIN · HEAP

    Last Stone Weight in Kotlin, explained

    Problem overview

    You have an IntArray of stone weights. On each turn you take the two heaviest stones and smash them together. If they weigh the same, both disappear. If not, the lighter one disappears and the heavier one shrinks to the difference. The game ends when at most one stone remains, and you return its weight, or 0 if none is left.

    For [2, 7, 4, 1, 8, 1] the turns go like this:

    1. 8 and 7 smash, leaving 1: the stones are 2, 4, 1, 1, 1.
    2. 4 and 2 smash, leaving 2: the stones are 2, 1, 1, 1.
    3. 2 and 1 smash, leaving 1: the stones are 1, 1, 1.
    4. 1 and 1 smash and vanish: one stone of weight 1 is left.

    The answer is 1.

    Approach

    Every turn needs "the heaviest stone", and the new stone that a smash leaves behind must go back among the others. Keeping the stones in a sorted array would work, but inserting into it each time costs O(n). Re-sorting every turn costs even more.

    A max heap fits the job. It gives you the greatest element in O(log n) and lets you insert a new element in O(log n). The algorithm is a direct simulation:

    1. Put all the stones in a max heap.
    2. While at least two stones remain, take the two heaviest.
    3. If they differ, put their difference back.
    4. When the loop ends, the heap is empty (answer 0) or holds one stone (that is the answer).

    Step by step

    1. Create the heap. Kotlin on the JVM uses Java's PriorityQueue, which is a min heap by default. The comparator compareByDescending { it } reverses the order, so the heaviest stone comes out first. Name it heap.
    2. Load the stones. stones is an IntArray, and addAll expects a collection, so stones.toList() converts the array first.
    3. Loop while two stones remain. The condition heap.size >= 2 means a smash is possible. When the loop ends, at most one stone is left.
    4. Take the heaviest. val x = heap.poll() removes and returns the top.
    5. Take the second heaviest. val y = heap.poll() comes out next, so x >= y.
    6. Push the remainder. If x != y, the new stone weighs x - y and goes back with heap.add. If they are equal, both vanish and nothing goes back.
    7. Empty heap. After the loop, if heap.size == 0, return 0.
    8. One stone left. Otherwise return heap.peek(), which reads the only stone without removing it.

    The solution in Kotlin

    class Solution {
    fun lastStoneWeight(stones: IntArray): Int {
        val heap = PriorityQueue<Int>(compareByDescending { it })
        heap.addAll(stones.toList())
        while (heap.size >= 2) {
            val x = heap.poll()
            val y = heap.poll()
            if (x != y) {
                heap.add(x - y)
            }
        }
        if (heap.size == 0) {
            return 0
        }
        return heap.peek()
    }
}

    Kotlin details worth knowing:

    • PriorityQueue is a Java class and lives in java.util. A standalone Kotlin file that uses the short name needs import java.util.PriorityQueue at the top.
    • compareByDescending { it } builds a comparator from a selector, here the number itself. Kotlin infers the type from PriorityQueue<Int>.
    • poll() and peek() are Java methods whose result Kotlin treats as a platform type, so you can assign them to an Int directly. Java's versions return null for an empty queue, so a call on an empty heap would fail when the value is used as an Int. Here every poll() is guarded by heap.size >= 2.
    • x != y compares the numbers by value. In Kotlin != calls equals, so it also works for boxed integers. In Java, comparing two Integer objects with != would compare references.
    • stones.toList() makes a copy, so the caller's array is not changed.

    Complexity

    Time: O(n log n). Loading the stones costs at most O(n log n) with repeated insertion. Each turn removes two stones and adds at most one, so the heap shrinks by at least one per turn and the loop runs fewer than n times. Each turn does two polls and at most one add, each O(log n).

    Space: O(n). The heap holds all the stones, and toList() creates one temporary copy of the same size.

    Edge cases and common mistakes

    • One stone. The loop never runs, and the answer is that stone's weight. For [5] the function returns 5.
    • Everything cancels. For [3, 3] the two stones destroy each other, the heap is empty and the function returns 0.
    • Forgetting the equal case. Pushing x - y when x == y would add a stone of weight 0 and could give a wrong answer.
    • Using the default heap. With a min heap you would smash the two lightest stones, which is a different game.
    • Reading the answer from an empty heap. heap.peek() on an empty heap gives null, not 0, so the size check must come first.

    Testing

    • [2, 7, 4, 1, 8, 1] returns 1.
    • [3, 3] returns 0.
    • [5] returns 5.
    Keep going with KotlinAbout 340 step-by-step lessons, from your first line of code to coding interviews. Free to start.
    Start learning free