Last Stone Weight

NeetCode: Last Stone Weight · LeetCode: 1046. Last Stone Weight · Difficulty: Easy

Each turn, take the two heaviest stones and smash them. If their weights are equal both vanish; otherwise the lighter one vanishes and the heavier becomes y - x. Return the weight of the last stone, or 0 if none remain.

Example

Text

Input:  stones = [2, 3, 6, 2, 4]    Output: 1
Input:  stones = [1, 2]             Output: 1

Brute force: sort every round

The code copies the stones into an ArrayList<Integer>. list.remove(int index) removes by position; removing the last position is O(1). The values y and x are int, so y != x compares numbers, not Integer references.

Java

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

class Solution {
    public int lastStoneWeight(int[] stones) {
        List<Integer> list = new ArrayList<>();
        for (int s : stones) {
            list.add(s);
        }
        while (list.size() > 1) {
            Collections.sort(list);
            int y = list.remove(list.size() - 1);
            int x = list.remove(list.size() - 1);
            if (y != x) {
                list.add(y - x);
            }
        }
        return list.isEmpty() ? 0 : list.get(0);
    }
}

public class Main {
    public static void main(String[] args) {
        Solution s = new Solution();
        System.out.println(s.lastStoneWeight(new int[]{2, 3, 6, 2, 4}));
        System.out.println(s.lastStoneWeight(new int[]{1, 2}));
        System.out.println(s.lastStoneWeight(new int[]{5, 5}));
    }
}

Output

1
1
0

Time: O(n² log n). Space: O(n).

Optimal: a max-heap

The operation is always "take the two largest, put one value back", which is exactly what a max-heap is for: O(log n) per poll and add. Java's PriorityQueue is a min-heap by default, so Collections.reverseOrder() turns it into a max-heap.

Java

import java.util.Collections;
import java.util.PriorityQueue;

class Solution {
    public int lastStoneWeight(int[] stones) {
        PriorityQueue<Integer> heap = new PriorityQueue<>(Collections.reverseOrder()); // max-heap
        for (int s : stones) {
            heap.add(s);
        }
        while (heap.size() > 1) {
            int y = heap.poll();
            int x = heap.poll();
            if (y != x) {
                heap.add(y - x);
            }
        }
        return heap.isEmpty() ? 0 : heap.poll();
    }
}

public class Main {
    public static void main(String[] args) {
        Solution s = new Solution();
        System.out.println(s.lastStoneWeight(new int[]{2, 3, 6, 2, 4}));
        System.out.println(s.lastStoneWeight(new int[]{1, 2}));
        System.out.println(s.lastStoneWeight(new int[]{5, 5}));
    }
}

Output

1
1
0

Time: O(n log n). Space: O(n).