Problem statement

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

Solution

There is straightforward O(n^3) solution, but it gives TLE after they added more tests. The idea is dp(i, j) be the answer for subproblem [i, j].

Complexity

Time is O(n^3), space is O(n^2).

Code

class Solution:
    def stoneGameV(self, A):
        acc = [0] + list(accumulate(A))
        
        @lru_cache(None)
        def dp(i, j):
            if i == j: return 0
            S = acc[j+1] - acc[i]
            ans = 0
            for k in range(i, j):
                lft = acc[k+1] - acc[i]
                rgh = S - lft
                if 2*lft <= S: ans = max(ans, lft + dp(i, k))
                if 2*rgh <= S: ans = max(ans, rgh + dp(k+1, j))     
            return ans
        
        return dp(0, len(A) - 1)

Solution 2

There are solutions with O(n^2 log n) and even with O(n^2) complexities, here is O(n^2) solution. The idea is to use so-called Knuth optimization: actually we do not need to traverse all k in previous solutions, but only edge cases: the moment where the sum of elements in left part becomes bigger than sum of right part. I will leave like this without explanation, I think it is a bit out of scope.

Code

class Solution:
    def stoneGameV(self, A):
        n = len(A)
        dp = [[0] * n for _ in range(n)]
        mx = [[0] * n for _ in range(n)]
        for i in range(n): mx[i][i] = A[i]
        
        for j in range(1, n):
            mid, sm, right = j, A[j], 0
            for i in range(j-1, -1, -1):
                sm += A[i]
                while (right + A[mid])*2 <= sm:
                    right += A[mid]
                    mid -= 1
                if right*2 == sm: dp[i][j] = mx[i][mid]
                if mid != i: dp[i][j] = max(dp[i][j], mx[i][mid - 1])       
                if mid != j: dp[i][j] = max(dp[i][j], mx[j][mid + 1])
                mx[i][j] = max(mx[i][j-1], dp[i][j] + sm)
                mx[j][i] = max(mx[j][i+1], dp[i][j] + sm)
    
        return dp[0][n-1]