三个无重叠子数组的最大和

标签: 数组 动态规划

难度: Hard

给你一个整数数组 nums 和一个整数 k ,找出三个长度为 k 、互不重叠、且全部数字和(3 * k 项)最大的子数组,并返回这三个子数组。

以下标的数组形式返回结果,数组中的每一项分别指示每个子数组的起始位置(下标从 0 开始)。如果有多个结果,返回字典序最小的一个。

示例 1:

输入:nums = [1,2,1,2,6,7,5,1], k = 2
输出:[0,3,5]
解释:子数组 [1, 2], [2, 6], [7, 5] 对应的起始下标为 [0, 3, 5]。
也可以取 [2, 1], 但是结果 [1, 3, 5] 在字典序上更大。

示例 2:

输入:nums = [1,2,1,2,1,2,1,2,1], k = 2
输出:[0,2,4]

提示:

  • 1 <= nums.length <= 2 * 104
  • 1 <= nums[i] < 216
  • 1 <= k <= floor(nums.length / 3)

Submission

运行时间: 48 ms

内存: 18.0 MB

class Solution:
    def maxSumOfThreeSubarrays(self, nums: List[int], k: int) -> List[int]:
        s = s1 = s2 = s3 = 0
        mx1 = mx12 = 0
        idx1, idx12 = 0, ()
        ans = []
        for i in range(k * 2, len(nums)):
            s1 += nums[i - k * 2]
            s2 += nums[i - k]
            s3 += nums[i]
            if i >= k * 3 - 1:
                if s1 > mx1:
                    mx1 = s1
                    idx1 = i - k * 3 + 1
                if mx1 + s2 > mx12:
                    mx12 = mx1 + s2
                    idx12 = (idx1, i - k * 2 + 1)
                if mx12 + s3 > s:
                    s = mx12 + s3
                    ans = [*idx12, i - k + 1]
                s1 -= nums[i - k * 3 + 1]
                s2 -= nums[i - k * 2 + 1]
                s3 -= nums[i - k + 1]
        return ans

Explain

本题解采用滑动窗口和动态规划的思想。首先,使用三个变量s1, s2, s3分别维护三个长度为k的子数组的和。然后,从第k * 2个元素开始遍历数组,每次向右移动一个元素,更新三个子数组的和。同时,使用mx1和mx12分别记录遍历过程中第一个子数组和前两个子数组的最大和,以及对应的起始下标idx1和idx12。最后,每次遍历时比较当前三个子数组的总和是否大于之前记录的最大值s,如果大于,则更新s和结果ans。

时间复杂度: O(n)

空间复杂度: O(1)

class Solution:
    def maxSumOfThreeSubarrays(self, nums: List[int], k: int) -> List[int]:
        s = s1 = s2 = s3 = 0  # 初始化三个子数组的和以及总和
        mx1 = mx12 = 0  # 初始化前一个和前两个子数组的最大和
        idx1, idx12 = 0, ()  # 初始化最大和对应的起始下标
        ans = []
        for i in range(k * 2, len(nums)):
            s1 += nums[i - k * 2]  # 更新第一个子数组的和
            s2 += nums[i - k]  # 更新第二个子数组的和
            s3 += nums[i]  # 更新第三个子数组的和
            if i >= k * 3 - 1:
                if s1 > mx1:
                    mx1 = s1
                    idx1 = i - k * 3 + 1
                if mx1 + s2 > mx12:
                    mx12 = mx1 + s2
                    idx12 = (idx1, i - k * 2 + 1)
                if mx12 + s3 > s:
                    s = mx12 + s3
                    ans = [*idx12, i - k + 1]
                s1 -= nums[i - k * 3 + 1]  # 移除第一个子数组最左边的元素
                s2 -= nums[i - k * 2 + 1]  # 移除第二个子数组最左边的元素
                s3 -= nums[i - k + 1]  # 移除第三个子数组最左边的元素
        return ans

Explore

在代码中,通过固定每个子数组的长度为k,并且按照严格的位置偏移来更新每个子数组的和来保证子数组互不重叠。具体来说,第一个子数组从索引0开始,到k-1结束;第二个子数组从索引k开始,到2k-1结束;第三个子数组从索引2k开始,到3k-1结束。随着数组的遍历,每个子数组都严格按照k个元素的长度向右滑动,通过这种固定长度和起始位置的偏移量控制,保证了子数组之间不会重叠。

在代码中,每次更新子数组和时,都会加上新进入窗口的元素,并移除窗口最左侧的元素。这种操作确保了每次窗口(或子数组)的和都是基于最新的k个连续元素。例如,s1更新时加上nums[i - k * 2]并减去nums[i - k * 3 + 1],这样s1始终保持为最新的从i - 2k + 1到i - k * 2的子数组的和。

确实可能出现多个总和相同但起始位置不同的情况。在实现中,通过只在找到更大的总和时更新结果来确保只保存最大值。如果遇到相同的总和,由于我们不更新结果,所以自然保留了最先找到的(即字典序最小的)结果。这是因为数组是按索引顺序遍历的,第一次遇到的相同最大总和的起始位置自然是最小的。

动态规划在本题中主要体现在逐步构建最优解,即不断更新记录最大子数组和的过程中。使用mx1来记录到当前位置为止,最大的第一个子数组的和。进一步,使用mx12来记录到当前位置为止,最大的前两个子数组的累加和。这种从前向后逐步优化求解的方法,利用以往的计算结果来简化未来的计算,是动态规划的典型应用。每步更新mx1和mx12都是基于之前的状态进行的,确保了每步都是基于最优子结构的决策。