Stone Game II

NeetCode: Stone Game II · LeetCode: 1140. Stone Game II · Difficulty: Medium

Alice and Bob play with piles piles[0..n] in a row, Alice first. Initially M = 1. On a turn, the player takes all the stones of the first X remaining piles, where 1 <= X <= 2M, and then sets M = max(M, X). The game ends when no piles remain. Both play optimally to maximise their own total. Return the maximum number of stones Alice can get.

Example

Text

Input:  piles = [2, 7, 9, 4, 4]         Output: 10
Input:  piles = [1, 2, 3, 4, 5, 100]    Output: 104

Brute force: plain recursion

Key insight: the game state is fully described by the index i of the first remaining pile and the current M. Let best(i, M) be the most stones the player to move can collect from piles i... The two players split the remaining stones, so if the mover takes X piles, the opponent then collects best(i + X, max(M, X)) and the mover keeps suffix[i] - best(i + X, max(M, X)). The mover picks the X that maximises this. If 2M covers all remaining piles, taking them all is best: suffix[i].

Without caching, every state branches into up to 2M children and identical (i, M) pairs are recomputed, so the cost is exponential in n.

Java

class Solution {
    public int stoneGameII(int[] piles) {
        int n = piles.length;
        int[] suffix = new int[n + 1];
        for (int i = n - 1; i >= 0; i--) suffix[i] = suffix[i + 1] + piles[i];
        return best(suffix, 0, 1);
    }

    private int best(int[] suffix, int i, int m) {
        int n = suffix.length - 1;
        if (i >= n) return 0;
        if (i + 2 * m >= n) return suffix[i];
        int result = 0;
        for (int x = 1; x <= 2 * m; x++) {
            result = Math.max(result, suffix[i] - best(suffix, i + x, Math.max(m, x)));
        }
        return result;
    }
}

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

Output

Time: exponential in n (each call spawns up to 2M calls, and states repeat). Space: O(n) recursion depth.

Better: memoise on (i, M)

There are n values of i, and M never usefully exceeds n, so there are at most n^2 distinct states. Each state loops over at most 2M <= 2n choices, giving O(n^3) in total. The code is the brute force plus a memo table. The table holds 0 for "not computed yet", which is safe because a state with piles left always has a result of at least 1 (piles hold at least one stone).

The suffix sum makes one function enough: the opponent's stones are best(i + x, max(M, x)), so the mover's are everything left minus that. No second function or player flag is needed.

Java

class Solution {
    public int stoneGameII(int[] piles) {
        int n = piles.length;
        int[] suffix = new int[n + 1];
        for (int i = n - 1; i >= 0; i--) suffix[i] = suffix[i + 1] + piles[i];
        int[][] memo = new int[n + 1][n + 1]; // 0 means "not computed yet"
        return best(suffix, memo, 0, 1);
    }

    private int best(int[] suffix, int[][] memo, int i, int m) {
        int n = suffix.length - 1;
        if (i >= n) return 0;
        if (i + 2 * m >= n) return suffix[i]; // can take everything left
        if (memo[i][m] != 0) return memo[i][m];
        int result = 0;
        for (int x = 1; x <= 2 * m; x++) {
            result = Math.max(result, suffix[i] - best(suffix, memo, i + x, Math.max(m, x)));
        }
        memo[i][m] = result;
        return result;
    }
}

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

Output

Time: O(n^3). Space: O(n^2) table plus O(n) recursion.

Optimal: bottom-up table

Let dp[i][m] be the most stones the player to move can collect from piles i.. when the current limit is m. Fill i from n - 1 down to 0, so every dp[i + x][..] is ready. For each m in 1..=n: if i + 2m >= n, dp[i][m] = suffix[i]; otherwise dp[i][m] = max over x in 1..=2m of suffix[i] - dp[i + x][max(m, x)]. The answer is dp[0][1].

Java

class Solution {
    public int stoneGameII(int[] piles) {
        int n = piles.length;
        if (n == 0) return 0;
        int[] suffix = new int[n + 1];
        for (int i = n - 1; i >= 0; i--) suffix[i] = suffix[i + 1] + piles[i];
        // dp[i][m]: best total for the player to move from piles i.. with limit m
        int[][] dp = new int[n + 1][n + 1];
        for (int i = n - 1; i >= 0; i--) {
            for (int m = 1; m <= n; m++) {
                if (i + 2 * m >= n) {
                    dp[i][m] = suffix[i]; // can take everything left
                } else {
                    int result = 0;
                    for (int x = 1; x <= 2 * m; x++) {
                        result = Math.max(result, suffix[i] - dp[i + x][Math.max(m, x)]);
                    }
                    dp[i][m] = result;
                }
            }
        }
        return dp[0][1];
    }
}

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

Output

Time: O(n^3). Space: O(n^2). The table cannot shrink to one row because state i reads many later indices and m values.

Edge cases

One pile: 2M >= 1 covers it, so Alice takes it.
The "take everything" shortcut is correct because taking all remaining stones is always at least as good when allowed.
M is capped at n, since a larger limit gives no extra choices.
With the problem's limits (at most 100 piles of at most 10,000 stones), every total fits in int.

How to explain it in an interview

Start with recursion over every X for both players; it is exponential because (i, M) pairs repeat.
Key insight: the state is (i, M), and with suffix sums the opponent's optimal take is subtracted from what remains, so one function best(i, M) covers both players.
Steps: build suffix[i]; if i + 2m >= n, return suffix[i]; otherwise take the max over x in 1..=2m of suffix[i] - best(i + x, max(m, x)); memoise on (i, m) or fill a table from the back.
Complexity: O(n^3) time, O(n^2) space. Alice's answer is best(0, 1).