Leetcode 0689 Maximum Sum of 3 Non-Overlapping Subarrays

Problem statement

https://leetcode.com/problems/maximum-sum-of-3-non-overlapping-subarrays/

Solution

Let us reformulate our problem: evaluate in O(n) sums of subarrays of size k and call this array windows. Let dp[i][j] be an answer for the question: what is the maximum sum of i non-overlapping windows if we allowed to use only nums[:j]. And index[i][j] is used for recording for ith interval, what are the current start index for first j numbers that made up the max sum. In the end we need to backtrack to form solution.

Complexity

Time and space complexity is O(mn) = O(n), because in this problem m = 3.

Code 1

class Solution:
    def maxSumOfThreeSubarrays(self, nums, k):
        m, n = 3, len(nums)
        dp = [[0]*(n+1) for _ in range(m+1)] 
        index = [[0]*(n+1) for _ in range(m+1)] 
        
        acc = [0] + list(accumulate(nums))
        windows = [acc[i+k] - acc[i] for i in range(n-k+1)] 

        for i in range(1, m + 1):
            for j in range(k - 1, n):
                tmpMax = windows[j+1-k] + dp[i-1][j-k+1]
                if tmpMax > dp[i][j]:
                    dp[i][j+1] = tmpMax
                    index[i][j+1] = j - k + 1
                else:
                    dp[i][j+1] = dp[i][j]
                    index[i][j+1] = index[i][j]
                    
        prev, ret = n, []
        for i in range(m, 0, -1):
            ret.append(index[i][prev])
            prev = ret[-1]
            
        return ret[::-1]

Code 2

It can be written in shorter way (but more difficult to understand) if we combine dp and index tables in one table of tuples. Also we need to use negative index trick to make sure we choose the smallest index on each step.

class Solution:
    def maxSumOfThreeSubarrays(self, nums, k):
        m, n = 3, len(nums)
        dp = [[(0, 0)]*(n+1) for _ in range(m+1)]  
        
        cum_sum = [0] + list(accumulate(nums))
        Windows = [cum_sum[i+k] - cum_sum[i] for i in range(n-k+1)] 

        for i in range(1, m + 1):
            for j in range(k - 1, n):
                tmpMax = Windows[j+1-k] + dp[i-1][j-k+1][0]
                dp[i][j+1] = max((tmpMax, -j + k - 1), dp[i][j])
                  
        prev, ret = n, []
        for i in range(m, 0, -1):
            ret.append(-dp[i][prev][1])
            prev = ret[-1]
            
        return ret[::-1]