Leetcode Daily Question 20/08/2024 - 1140. Stone Game II

Problem Description:

Alice and Bob continue their games with piles of stones. There are a number of piles arranged in a row, and each pile has a positive integer number of stones piles[i]. The objective of the game is to end with the most stones. Alice and Bob take turns, with Alice starting first. Initially, M = 1. On each player's turn, that player can take all the stones in the first X remaining piles, where 1 <= X <= 2M. Then, we set M = max(M, X). The game continues until all the stones have been taken. Assuming Alice and Bob play optimally, return the maximum number of stones Alice can get.

from typing import List
from functools import lru_cache
class Solution:
  def stoneGamesII(self, piles:List[int]) -> int:
    n = len(piles)
    suffix = [0] * (n+1)
    for i in range(n-1,-1,-1):
      suffix[i] = suffix[i+1] + piles[i]
      
    @lru_cache(None)
    def dp(i,m):
      if i + 2*m >= n:
        return suffix[i]
      
      minScore = float('inf')
      
      for x in range(1, 2*m+1):
        minScore = min(minScore, dp(i+x, max(m,x)))
      return suffix[i] - minScore
    
    return dp(0,1)

Explanation:

My first instinct was to use dp to determine the best strategy to play the game. Let's first explore how we can do this then we shall explore a memorisation method to speed up the dp.

First, let's compute the sum of piles from i to the end of the array for convenience. This way we can immediately know the sum of the remaining piles to be taken to determine the optimum strategy.

def stoneGameII(piles): 
  n = len(piles)
  suffix = [0] * (n+1) #we need 1 extra cell to make the traversal work
  for i in range(n-1,-1,-1):
    suffix[i] = suffix[i+1] + piles[i] #previous cumulative suffix sum + current pile

When it comes to doing the dp, we use the dp structure dp[i][m] where m = M in the question of possible moves to make and i represent the index from the end of the array.

Our strategy is to try taking every pile from 1 to 2*m to decide which path is best to take.

Since we know our opponent also always choose the optimum strategy we can predict what they would do.

def stoneGameII(piles): 
  n = len(piles)
  suffix = [0] * (n+1) #we need 1 extra cell to make the traversal work
  for i in range(n-1,-1,-1):
    suffix[i] = suffix[i+1] + piles[i] #previous cumulative suffix sum + current pile
	
  dp = [[0]*(n+1) for _ in range(n+1)]
  for i in range(n-1,-1,-1):
    for m in range(1, n+1):
      max_stones = 0
      for x in range(1, 2*m + 1):
        if i + x <= n:
        	max_stones = max(max_stones, suffix[i] - dp[i+x][max(m,x)])
      dp[i][m] = max_stones
  return dp[0][1]

This gives a time complexity of O(n^3) and space complexity of O(n^2) since we need to go through every possibility and store results in a list of list.

Now we can do the same dp, but with memorisation to speed up the process. We do this using lru_cache from the functools module in python.

We start as we did, calculating the cumulative sum from each i to the end of piles.

def stoneGameII(piles):
  n = len(piles)
  suffix = [0] * (n+1)
  for i in range(n-1,-1,-1):
    suffix[i] = suffix[i+1] + piles[i]
    
  @lru_cache(None)
  def dp(i,m):
    if i + 2*m >= n: #if the maximum moves we can make reach outside the piles, we take the rest of the piles.
      return suffix[i] 
    
   	minScore = float('inf')
    
    for x in range(1, 2*m+1): #try every possible path
      minScore = min(minScore, dp(i+x, max(m,x)))
    return suffix[i] - minScore
  
  return dp(0,1)

This has maximum time complexity of O(n^3). This is because there are O(n^2) possible states (due to i and m), and for each state, we might do up to O(n) work due to the loop over x.

Space complexity of O(n^2).

The reason why caching is so much faster is because we memorise every path and avoid repeated calculations.

The full solution can be found near the top of the page.