Problem statement

https://leetcode.com/problems/stone-game-ii/

Solution

Let dp(i, M) be the maximum gain the if the first player have piles left [i, ... ] and he can take >= M of them. Then how many stones he can get? We can look at place j from i + 1 to i + 2*M + 1 and we can have sum(A[i:]) - gain of the second player, which can be evaluated as dp(j, max(M, j - i)).

Complexity

It is O(n^3) for time and O(n^2) for space.

Code

class Solution:
    def stoneGameII(self, A):
        n, acc = len(A), [0] + list(accumulate(A))

        @lru_cache(None)
        def dp(i, M):
            res = 0
            for j in range(i + 1, min(2 * M + i + 1, n + 1)):
                res = max(res, acc[-1] - acc[i] - dp(j, max(M, j - i)))
            return res

        return dp(0, 1)