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:
-
C++
-
Python
#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) \).
- The function
- 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 thesubArraySums
vector. - Each state computation involves constant-time operations, resulting in \( O(3 \times n) = O(n) \) overall.
- The function
- Solution Reconstruction:
- The function
solve
iterates through thesubArraySums
vector, making at most \( O(n) \) iterations.
- The function
- Overall Time Complexity:
The total time complexity is: \( O(n) \).
Space Complexity
- Subarray Sums:
- The
subArraySums
vector requires \( O(n) \) space.
- The
- Dynamic Programming Table:
- The
dp
table stores up to \( O(3 \times n) \) states, requiring \( O(n) \) space.
- The
- 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) \).