1140. Stone Game II - cocoder39/coco39_LC GitHub Wiki

1140. Stone Game II

dfs(i, M) returns the max num of stones current player can get if the remaining piles are [i ... n-1]

Time complexity: upper bound of M is n, so in total there are n^2 subproblems. Each subproblem takes 2M operations so total time complexity is O(n ^ 3)

class Solution:
    def stoneGameII(self, piles: List[int]) -> int:
        n = len(piles)
        mem = collections.defaultdict(dict)
        
        total = [0] * n 
        for i in range(n-1, -1, -1):
            if i == n-1:
                total[i] = piles[i]
            else:
                total[i] = total[i+1] + piles[i]
        
        def helper(i, M):
            if i == len(piles):
                return 0
            
            if i in mem and M in mem[i]:
                return mem[i][M]
            
            if i + 2*M >= len(piles):
                return total[i]
            
            res = 0
            for x in range(1, 2*M+1):
                j = i + x - 1
                res = max(res, total[i] - helper(j+1, max(M, x)))
                
            mem[i][M] = res
            return res
        
        return helper(0, 1)