689. Maximum Sum of 3 Non-Overlapping Subarrays

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 * 104
  • 1 <= nums[i] < 216
  • 1 <= k <= floor(nums.length / 3)

Approach 01:

#include <vector>
#include <limits>
#include <algorithm>
using namespace std;

class Solution {
public:
    vector<int> calculateSubArraySums(vector<int>& nums, int subArraySize) {
        int currentSum = 0;
        vector<int> subArraySums;

        for (int i = 0; i < nums.size(); ++i) {
            currentSum += nums[i];

            // Start forming subarrays after the first subArraySize elements
            if (i >= subArraySize - 1) {
                subArraySums.push_back(currentSum);
                currentSum -= nums[i - subArraySize + 1];
            }
        }

        return subArraySums;
    }

    int calculateMaxSum(int remainingCount, int startIndex, int subArraySize, vector<int>& subArraySums, vector<vector<int>>& dp) {
        if (startIndex >= subArraySums.size() || remainingCount == 0) {
            return 0;
        }

        if (dp[remainingCount][startIndex] != -1) {
            return dp[remainingCount][startIndex];
        }

        // Include current subarray
        int includeCurrent = subArraySums[startIndex] + calculateMaxSum(remainingCount - 1, startIndex + subArraySize, subArraySize, subArraySums, dp);

        // Skip current subarray
        int excludeCurrent = calculateMaxSum(remainingCount, startIndex + 1, subArraySize, subArraySums, dp);

        dp[remainingCount][startIndex] = max(includeCurrent, excludeCurrent);
        return dp[remainingCount][startIndex];
    }

    void solve(int subArraySize, int remainingCount, int currentIndex, vector<int>& resultIndices, vector<int>& subArraySums, vector<vector<int>>& dp) {
        while (remainingCount > 0 && currentIndex < subArraySums.size()) {
            int includeCurrent = subArraySums[currentIndex] + calculateMaxSum(remainingCount - 1, currentIndex + subArraySize, subArraySize, subArraySums, dp);
            int excludeCurrent = calculateMaxSum(remainingCount, currentIndex + 1, subArraySize, subArraySums, dp);

            if (includeCurrent >= excludeCurrent) {
                resultIndices.push_back(currentIndex);
                remainingCount--;
                currentIndex += subArraySize;
            } else {
                currentIndex++;
            }
        }
    }

    vector<int> maxSumOfThreeSubarrays(vector<int>& nums, int subArraySize) {
        vector<int> resultIndices;
        vector<int> subArraySums = calculateSubArraySums(nums, subArraySize);
        int n = subArraySums.size();
        vector<vector<int>> dp(4, vector<int>(n, -1)); // 4 rows for up to 3 subarrays + base case

        solve(subArraySize, 3, 0, resultIndices, subArraySums, dp);
        return resultIndices;
    }
};
from typing import List

class Solution:
    def calculateSubArraySums(self, nums: List[int], subArraySize: int) -> List[int]:
        currentSum = 0
        subArraySums = []

        for i in range(len(nums)):
            currentSum += nums[i]

            # Start forming subarrays after the first subArraySize elements
            if i >= subArraySize - 1:
                subArraySums.append(currentSum)
                currentSum -= nums[i - subArraySize + 1]

        return subArraySums

    def calculateMaxSum(self, remainingCount: int, startIndex: int, subArraySize: int, subArraySums: List[int], dp: List[List[int]]) -> int:
        if startIndex >= len(subArraySums) or remainingCount == 0:
            return 0

        if dp[remainingCount][startIndex] != -1:
            return dp[remainingCount][startIndex]

        # Include current subarray
        includeCurrent = subArraySums[startIndex] + self.calculateMaxSum(remainingCount - 1, startIndex + subArraySize, subArraySize, subArraySums, dp)

        # Skip current subarray
        excludeCurrent = self.calculateMaxSum(remainingCount, startIndex + 1, subArraySize, subArraySums, dp)

        dp[remainingCount][startIndex] = max(includeCurrent, excludeCurrent)
        return dp[remainingCount][startIndex]

    def solve(self, subArraySize: int, remainingCount: int, currentIndex: int, resultIndices: List[int], subArraySums: List[int], dp: List[List[int]]):
        while remainingCount > 0 and currentIndex < len(subArraySums):
            includeCurrent = subArraySums[currentIndex] + self.calculateMaxSum(remainingCount - 1, currentIndex + subArraySize, subArraySize, subArraySums, dp)
            excludeCurrent = self.calculateMaxSum(remainingCount, currentIndex + 1, subArraySize, subArraySums, dp)

            if includeCurrent >= excludeCurrent:
                resultIndices.append(currentIndex)
                remainingCount -= 1
                currentIndex += subArraySize
            else:
                currentIndex += 1

    def maxSumOfThreeSubarrays(self, nums: List[int], subArraySize: int) -> List[int]:
        resultIndices = []
        subArraySums = self.calculateSubArraySums(nums, subArraySize)
        n = len(subArraySums)
        dp = [[-1] * n for _ in range(4)]  # 4 rows for up to 3 subarrays + base case

        self.solve(subArraySize, 3, 0, resultIndices, subArraySums, dp)
        return resultIndices

Time Complexity

  • Subarray Sum Calculation:
    • The function calculateSubArraySums iterates through the array of size \( n \) once to compute the sum of all possible subarrays of size \( k \).
    • This step takes \( O(n) \).
  • Dynamic Programming Table Calculation:
    • The function calculateMaxSum uses recursion with memoization. It computes values for up to \( O(3 \times n) \) states, where \( 3 \) is the maximum number of subarrays, and \( n \) is the size of the subArraySums vector.
    • Each state computation involves constant-time operations, resulting in \( O(3 \times n) = O(n) \) overall.
  • Solution Reconstruction:
    • The function solve iterates through the subArraySums vector, making at most \( O(n) \) iterations.
  • Overall Time Complexity:

    The total time complexity is: \( O(n) \).

Space Complexity

  • Subarray Sums:
    • The subArraySums vector requires \( O(n) \) space.
  • Dynamic Programming Table:
    • The dp table stores up to \( O(3 \times n) \) states, requiring \( O(n) \) space.
  • Auxiliary Variables:
    • Additional variables, such as loop counters and intermediate values, require \( O(1) \) space.
  • Overall Space Complexity:

    The total space complexity is: \( O(n) \).

Leave a Comment

Your email address will not be published. Required fields are marked *

Scroll to Top