LeetCode 689: Maximum Sum of 3 Non-Overlapping Subarrays

Problem Description

Given an integer array nums and an integer k, find three non-overlapping subarrays of length k with maximum sum and return them.

Return the result as a list of indices representing the starting position of each interval (0-indexed). If there are multiple answers, return the lexicographically smallest one.

Example 1:

Input: nums = [1,2,1,2,6,7,5,1], k = 2
Output: [0,3,5]
Explanation: Subarrays [1, 2], [2, 6], [7, 5] correspond to the starting indices [0, 3, 5].
We could have also taken [2, 1], but an answer of [1, 3, 5] would be lexicographically larger.

Example 2:

Input: nums = [1,2,1,2,1,2,1,2,1], k = 2
Output: [0,2,4]

Constraints:

1 <= nums.length <= 2 * 10⁴
1 <= nums[i] < 2¹⁶
1 <= k <= floor(nums.length / 3)

Difficulty: Hard

Tags: array, dynamic programming

Rating: 94.35%

Solution

Here’s my Python solution to this problem:

class Solution:
    def maxSumOfThreeSubarrays(self, nums: List[int], k: int) -> List[int]:
        n = len(nums)
        prefix_sum = [0] * (n + 1)
        for i in range(n):
            prefix_sum[i + 1] = prefix_sum[i] + nums[i]
        
        # Function to get sum of subarray starting at index i with length k
        def get_subarray_sum(i):
            return prefix_sum[i + k] - prefix_sum[i]
        
        # dp[i][j] represents the maximum sum possible using j + 1 subarrays 
        # considering the array up to index i
        dp = [[0] * 3 for _ in range(n)]
        
        # pos[i][j] stores the starting position of the last subarray used
        # in the optimal solution for dp[i][j]
        pos = [[0] * 3 for _ in range(n)]
        
        # Initialize first window sum
        dp[k-1][0] = get_subarray_sum(0)
        
        # Fill dp table for first array (j=0)
        for i in range(k, n):
            curr_sum = get_subarray_sum(i-k+1)
            if curr_sum > dp[i-1][0]:
                dp[i][0] = curr_sum
                pos[i][0] = i-k+1
            else:
                dp[i][0] = dp[i-1][0]
                pos[i][0] = pos[i-1][0]
        
        # Fill dp table for second and third arrays
        for j in range(1, 3):
            for i in range((j+1)*k-1, n):
                # Don't take current window
                dp[i][j] = dp[i-1][j]
                pos[i][j] = pos[i-1][j]
                
                # Take current window
                curr_sum = get_subarray_sum(i-k+1)
                prev_sum = dp[i-k][j-1]
                
                if curr_sum + prev_sum > dp[i][j]:
                    dp[i][j] = curr_sum + prev_sum
                    pos[i][j] = i-k+1
        
        # Backtrack to find positions
        result = []
        curr = n-1
        for j in range(2, -1, -1):
            result.insert(0, pos[curr][j])
            curr = pos[curr][j] - 1
            
        return result

Example Walkthrough

Let’s walk through how the solution works with the example:

Input: nums = [1,2,1,2,6,7,5,1], k = 2

First, we calculate the prefix sum array:

prefix_sum = [0,1,3,4,6,12,19,24,25]

We create our helper function get_subarray_sum(i) that uses the prefix sum to quickly calculate any subarray sum of length k=2:

# For example, to get sum of subarray [2,6] starting at index 3:
get_subarray_sum(3) = prefix_sum[5] - prefix_sum[3] = 12 - 4 = 8

We initialize our dp arrays:

dp = [[0,0,0] for _ in range(8)]   # Track maximum sums
pos = [[0,0,0] for _ in range(8)]  # Track positions

For j=0 (first subarray), we fill the first column of dp:

# At i=1: [1,2] = 3
dp[1][0] = 3
pos[1][0] = 0

# At i=2: max([1,2]=3 vs [2,1]=3)
dp[2][0] = 3
pos[2][0] = 0  # Keep lexicographically smaller

# Continue this process...

For j=1 (second subarray), we consider combinations:

# At i=3: [1,2] + [2,6] = 3 + 8 = 11
dp[3][1] = 11
pos[3][1] = 3

# Continue for remaining positions...

For j=2 (third subarray), we find:

# Best combination ends up being:
# [1,2] (at index 0) = 3
# [2,6] (at index 3) = 8
# [7,5] (at index 5) = 12
# Total sum = 23

Finally, we backtrack through our pos array to get the starting positions:

result = [0,3,5]

This gives us [0,3,5] as the answer, representing the positions of the three non-overlapping subarrays [1,2], [2,6], and [7,5] that give us the maximum possible sum of 23.

The solution is guaranteed to give us the lexicographically smallest answer because at each step, when we have a tie in sums, we keep the smaller position value.

Complexity Analysis

The solution has the following complexity characteristics:

Time Complexity: $O(n)$
Space Complexity: $O(n)$

Note: This is an automated analysis and may not capture all edge cases or specific algorithmic optimizations.