题目:
给你一个整数数组 nums 和一个整数 k ,找出三个长度为 k 、互不重叠、且 3 * k 项的和最大的子数组,并返回这三个子数组。
以下标的数组形式返回结果,数组中的每一项分别指示每个子数组的起始位置(下标从 0 开始)。如果有多个结果,返回字典序最小的一个。
思路:
前缀和+动态规划
解答:
class Solution:
def maxSumOfThreeSubarrays(self, nums: List[int], k: int) -> List[int]:
#前缀和+DP
#dp[i][j]:到数组第j个元素为止,前i个互不重叠的子数组的最大和
n=len(nums)
dp=[[0]*(n+1) for _ in range(4)]
#presum[i]:数组nums[:i+1]的前缀和
presum=[0]*(n+1)
for i in range(1,n+1):
presum[i]=presum[i-1]+nums[i-1]
#动态规划
for i in range(1,4):
for j in range(k*i,n+1):
#第j个元素(即:nums[j-1]) 要么不取,要么取
dp[i][j]=max(dp[i][j-1],dp[i-1][j-k]+presum[j]-presum[j-k])
#print(dp)
res=[0]*3
#回溯出字典序最小的具体方案
i,j,idx=3,n,2
while j>0:
#判断dp[i][j]从谁转移而来
while j!=0 and dp[i][j]==dp[i][j-1]:
j-=1
#若从dp[i][j-1]转化而来,则nums[j-1]未被选取,即nums[j-k:j]这段未被采纳
if dp[i][j-1]>dp[i-1][j-k]+presum[j]-presum[j-k]:
j-=1
#若从dp[i-1][j-k]转化而来,则nums[j-1]被选取,nums[j-k:j]该段被采纳
else:
j-=k
i-=1
res[idx]=j
idx-=1
return res



