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 * 1041 <= nums[i] < 2161 <= 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
calculateSubArraySumsiterates 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
calculateMaxSumuses 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 thesubArraySumsvector. - Each state computation involves constant-time operations, resulting in \( O(3 \times n) = O(n) \) overall.
- The function
- Solution Reconstruction:
- The function
solveiterates through thesubArraySumsvector, making at most \( O(n) \) iterations.
- The function
- Overall Time Complexity:
The total time complexity is: \( O(n) \).
Space Complexity
- Subarray Sums:
- The
subArraySumsvector requires \( O(n) \) space.
- The
- Dynamic Programming Table:
- The
dptable 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) \).