LeetCode 1140 - Stone Game II

The problem is a two-player turn-based game played on a row of stone piles. Each pile has a positive number of stones, and players take turns taking stones from the start of the remaining piles.

leetcodemediumarraymathdynamic-programmingprefix-sumgame-theory

LeetCode Problem 1140

Difficulty: 🟡 Medium
Topics: Array, Math, Dynamic Programming, Prefix Sum, Game Theory

Solution

Problem Understanding

The problem is a two-player turn-based game played on a row of stone piles. Each pile has a positive number of stones, and players take turns taking stones from the start of the remaining piles. On each turn, a player can take up to 2 * M piles, where M starts at 1 and is updated to max(M, X) after each move, with X being the number of piles taken in that turn. The objective for each player is to maximize the total number of stones they collect.

The input piles is an array of integers representing the number of stones in each pile. The output is the maximum number of stones Alice (the first player) can obtain if both Alice and Bob play optimally. The constraints (1 <= piles.length <= 100 and 1 <= piles[i] <= 10^4) suggest that we need an efficient solution because a naive recursive simulation of all possibilities could explode combinatorially.

Important edge cases include arrays with a single pile, arrays with two piles, and cases where the last piles have very high stone values, which could affect the optimal strategy significantly.

Approaches

The brute-force approach is to recursively simulate all possible choices Alice can make, for each choice recursively simulate Bob's optimal response, and continue until all piles are taken. This guarantees correctness because we explore all possible game paths, but it is extremely slow due to exponential branching at each step, specifically O(2^(n)) complexity in the worst case. This is infeasible for the input limit of 100 piles.

The key insight for an optimal solution is to use dynamic programming with memoization. Since the game is fully determined by the current index in the piles array and the current value of M, we can define a state (i, M) and store the result of optimal moves from this state. Using prefix sums, we can quickly calculate the sum of stones taken in a range without recomputation. This reduces redundant calculations and gives us a polynomial time solution.

Approach Time Complexity Space Complexity Notes
Brute Force O(2^n) O(n) Recursively simulate all choices without memoization. Exponential branching.
Optimal O(n^2 * n) = O(n^3) O(n^2) Use memoization with state (index, M) and prefix sums to compute max stones efficiently.

Algorithm Walkthrough

  1. Compute a suffix sum array suffix[i] where suffix[i] is the total stones from pile i to the end. This allows us to quickly compute the sum of stones remaining if a player takes the first X piles from index i.
  2. Define a recursive function dp(i, M) that returns the maximum number of stones the current player can collect starting from pile i with current M.
  3. In dp(i, M), consider all choices of taking X piles where 1 <= X <= min(2*M, n-i). For each choice, the current player gains suffix[i] - suffix[i+X] stones.
  4. The opponent will then optimally collect stones from the remaining piles, which is captured by dp(i+X, max(M, X)). To maximize the current player's gain, subtract the opponent's optimal gain from the total remaining stones.
  5. Memoize the result of dp(i, M) to avoid redundant computation. Return the memoized value if it exists.
  6. Call dp(0, 1) to get the maximum stones Alice can collect.

Why it works: The invariant is that at any state (i, M), the function dp(i, M) returns the best possible outcome for the current player assuming optimal play by both sides. Memoization ensures that each state is computed only once, so all overlapping subproblems are efficiently solved.

Python Solution

from typing import List

class Solution:
    def stoneGameII(self, piles: List[int]) -> int:
        n = len(piles)
        suffix = [0] * (n + 1)
        for i in range(n - 1, -1, -1):
            suffix[i] = piles[i] + suffix[i + 1]
        
        memo = {}
        
        def dp(i: int, M: int) -> int:
            if i >= n:
                return 0
            if (i, M) in memo:
                return memo[(i, M)]
            
            max_stones = 0
            for X in range(1, min(2 * M, n - i) + 1):
                opponent = dp(i + X, max(M, X))
                max_stones = max(max_stones, suffix[i] - opponent)
            
            memo[(i, M)] = max_stones
            return max_stones
        
        return dp(0, 1)

This implementation computes the suffix sum array to quickly get the total stones remaining. The recursive dp function iterates over possible moves X and recursively evaluates the opponent's optimal response, memoizing results to prevent repeated calculations. The final call dp(0, 1) represents Alice starting at the first pile with M = 1.

Go Solution

func stoneGameII(piles []int) int {
    n := len(piles)
    suffix := make([]int, n+1)
    for i := n-1; i >= 0; i-- {
        suffix[i] = piles[i] + suffix[i+1]
    }
    
    memo := make(map[[2]int]int)
    
    var dp func(int, int) int
    dp = func(i int, M int) int {
        if i >= n {
            return 0
        }
        key := [2]int{i, M}
        if val, ok := memo[key]; ok {
            return val
        }
        
        maxStones := 0
        for X := 1; X <= 2*M && i+X <= n; X++ {
            opponent := dp(i+X, max(M, X))
            maxStones = max(maxStones, suffix[i]-opponent)
        }
        memo[key] = maxStones
        return maxStones
    }
    
    return dp(0, 1)
}

func max(a, b int) int {
    if a > b { return a }
    return b
}

In Go, we handle memoization using a map with [2]int keys to represent (i, M). The logic is otherwise identical to the Python version. The max function is implemented explicitly since Go does not have a built-in one for integers.

Worked Examples

Example 1: piles = [2,7,9,4,4]

Step i M X chosen Stones taken Opponent optimal Current max
1 0 1 1 2 dp(1,1)=8 10
2 0 1 2 9 dp(2,2)=7 9

Alice chooses X=1 to maximize stones to 10.

Example 2: piles = [1,2,3,4,5,100]

Step i M X chosen Stones taken Opponent optimal Current max
1 0 1 1 1 dp(1,1)=104 104
2 0 1 2 3 dp(2,2)=100 104

Alice chooses X=1 for maximum stones 104.

Complexity Analysis

Measure Complexity Explanation
Time O(n^3) For each index i (n possibilities) and M (up to n), we try up to 2*M = O(n) choices, resulting in O(n^3).
Space O(n^2) Memoization table stores results for each (i, M) pair. Suffix sum array uses O(n).

The complexity is feasible given n <= 100.

Test Cases

# Provided examples
assert Solution().stoneGameII([2,7,9,4,4]) == 10  # Example 1
assert Solution().stoneGameII([1,2,3,4,5,100]) == 104  # Example 2

# Boundary cases
assert Solution().stoneGameII([1]) == 1  # Single pile
assert Solution().stoneGameII([1,1]) == 1  # Two equal piles

# Increasing sequence
assert Solution().stoneGameII([1,2,3,4,5]) == 9  # Optimal split

# Large numbers
assert Solution().stoneGameII([10000]*100) == 500000  # All piles same

# Mixed values
assert Solution().stoneGameII([5,3,4,5,6,2,1]) == 16  # Optimal selection pattern
Test Why
[2,7,9,4,4] Standard example, tests general strategy
[1,2,3,4,5,100] Tests impact of large value at end
[1] Minimum input edge case
[1,1]