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 smaller.
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)
Solution
Python3
class Solution:
def maxSumOfThreeSubarrays(self, nums: List[int], k: int) -> List[int]:
N = len(nums)
bestSeq = [0]
bestTwoSeq = [0, k]
bestThreeSeq = [0, k, k * 2]
seqSum = sum(nums[:k])
seqTwoSum = sum(nums[k : k * 2])
seqThreeSum = sum(nums[k * 2 : k * 3])
bestSeqSum = seqSum
bestTwoSum = seqSum + seqTwoSum
bestThreeSum = seqSum + seqTwoSum + seqThreeSum
seqIndex = 1
seqTwoIndex = k + 1
seqThreeIndex = 2 * k + 1
while seqThreeIndex <= N - k:
seqSum = seqSum - nums[seqIndex - 1] + nums[seqIndex + k - 1]
seqTwoSum = seqTwoSum - nums[seqTwoIndex - 1] + nums[seqTwoIndex + k - 1]
seqThreeSum = seqThreeSum - nums[seqThreeIndex - 1] + nums[seqThreeIndex + k - 1]
if seqSum > bestSeqSum:
bestSeqSum = seqSum
bestSeq = [seqIndex]
if bestSeqSum + seqTwoSum > bestTwoSum:
bestTwoSum = bestSeqSum + seqTwoSum
bestTwoSeq = bestSeq + [seqTwoIndex]
if bestTwoSum + seqThreeSum > bestThreeSum:
bestThreeSum = bestTwoSum + seqThreeSum
bestThreeSeq = bestTwoSeq + [seqThreeIndex]
seqIndex += 1
seqTwoIndex += 1
seqThreeIndex += 1
return bestThreeSeq