线性动态规划:最大子数组和问题的变种——至少包含k个元素的最大子数组和
题目描述
给定一个整数数组 nums 和一个整数 k (1 ≤ k ≤ n),请你找出数组中至少包含 k 个连续元素的子数组,使得这个子数组的和最大。返回这个最大的和。
示例1:
输入:nums = [1, -2, 3, 4, -1, 2, 1, -5, 4], k = 3
输出:10
解释:子数组 [3, 4, -1, 2, 1] 的和为 10,长度至少为 3。
示例2:
输入:nums = [-1, -2, -3, -4], k = 2
输出:-6
解释:子数组 [-3, -4] 的和为 -7,但子数组 [-1, -2] 的和为 -3,更小。然而,我们必须取至少长度为2的子数组。所以长度至少为2的最大和子数组是 [-1, -2] 或 [-3, -4],和为 -3?等等,我们计算一下:(-1) + (-2) = -3,(-3) + (-4) = -7。最大和是 -3。实际上,所有元素都是负数,我们不得不选择一个至少长度为k的子数组,那么最优解就是绝对值最小的那些负数组成的子数组。在这个例子中,长度为2的子数组最大和是 -3。
解题思路
这是一个经典的最大子数组和问题(Kadane算法)的变种,增加了子数组长度至少为 k 的限制。我们需要调整思路,不能简单地用滑动窗口(因为窗口是固定的k),而是要处理长度至少为k的情况。
核心难点
当子数组长度至少为 k 时,可能的子数组范围很广。直接枚举所有起始和结束位置是 O(n²),对于大数据会超时。我们需要用线性动态规划来优化。
关键观察
对于任意一个以位置 j 结束的子数组,如果我们想让它长度至少为 k,那么它的起始位置 i 必须满足:i ≤ j - k + 1(这样长度 ≥ k)。
换句话说,当我们固定结束位置 j 时,我们需要考虑从某个起始位置 i 到 j 的和,其中 i 的范围是 [0, j - k + 1]。
动态规划定义
我们可以定义两个辅助数组:
dp[j]:表示以位置j结束的、长度至少为k的子数组的最大和。maxEndingAt[j]:表示以位置j结束的普通最大子数组和(没有长度限制)。
但是这样定义 dp[j] 直接计算仍然需要枚举起始点。我们需要更聪明的递推。
优化递推公式
设 sum[i..j] 为从 i 到 j 的子数组和。
对于以 j 结束的长度至少为 k 的子数组,有两种情况:
- 子数组长度恰好为
k,即起始位置为j - k + 1。 - 子数组长度大于
k,即起始位置早于j - k + 1。
对于情况1,和就是 sum[j - k + 1 .. j],我们可以用前缀和快速计算。
对于情况2,这等价于:一个以 j-1 结束的长度至少为 k 的子数组,再添加上 nums[j]。
但是!这是不对的。因为以 j-1 结束的长度至少为 k 的子数组,加上 nums[j] 后长度至少为 k+1,仍然满足“长度至少为 k”。
所以我们有递推:
dp[j] = max( sum[j-k+1 .. j], dp[j-1] + nums[j] )
解释:
sum[j-k+1 .. j]:长度为k的子数组和。dp[j-1] + nums[j]:将至少长度为k且以j-1结尾的子数组向后扩展一位,得到的新子数组长度至少为k+1,也满足条件。
这样,我们就有了线性递推公式。
前缀和
为了快速计算任意子数组和,我们先用 O(n) 时间计算前缀和数组 prefix,其中 prefix[i] 表示前 i 个元素的和(prefix[0] = 0)。那么 sum[i..j] = prefix[j+1] - prefix[i]。
算法步骤
- 计算前缀和数组
prefix。 - 初始化
dp[k-1] = sum[0..k-1](即第一个长度为k的子数组的和)。 - 对于
j从k到n-1:- 当前长度为k的子数组和:
sum[j-k+1 .. j] = prefix[j+1] - prefix[j-k+1] dp[j] = max( prefix[j+1] - prefix[j-k+1], dp[j-1] + nums[j] )
- 当前长度为k的子数组和:
- 最终答案就是
max(dp[k-1], dp[k], ..., dp[n-1])。
例子演示
以 nums = [1, -2, 3, 4, -1, 2, 1, -5, 4], k = 3 为例:
- 前缀和 prefix = [0, 1, -1, 2, 6, 5, 7, 8, 3, 7]
- j=2 (第三个数):dp[2] = sum[0..2] = 1 + (-2) + 3 = 2。
- j=3:
sum[1..3] = prefix[4]-prefix[1] = 6 - 1 = 5
dp[2] + nums[3] = 2 + 4 = 6
dp[3] = max(5, 6) = 6 - j=4:
sum[2..4] = prefix[5]-prefix[2] = 5 - (-1) = 6
dp[3] + nums[4] = 6 + (-1) = 5
dp[4] = max(6, 5) = 6 - j=5:
sum[3..5] = prefix[6]-prefix[3] = 7 - 2 = 5
dp[4] + nums[5] = 6 + 2 = 8
dp[5] = max(5, 8) = 8 - j=6:
sum[4..6] = prefix[7]-prefix[4] = 8 - 6 = 2
dp[5] + nums[6] = 8 + 1 = 9
dp[6] = max(2, 9) = 9 - j=7:
sum[5..7] = prefix[8]-prefix[5] = 3 - 5 = -2
dp[6] + nums[7] = 9 + (-5) = 4
dp[7] = max(-2, 4) = 4 - j=8:
sum[6..8] = prefix[9]-prefix[6] = 7 - 7 = 0
dp[7] + nums[8] = 4 + 4 = 8
dp[8] = max(0, 8) = 8
所有 dp 值: [2, 6, 6, 8, 9, 4, 8],最大值是 10?等等,我们发现计算中最大值是 9,但示例答案是 10。哪里出错了?
检查错误
我们发现子数组 [3, 4, -1, 2, 1] 和是 3+4-1+2+1 = 9,不是 10。我们重新计算示例:
[3, 4, -1, 2, 1] = 9。那么示例中给出的 10 是怎么来的?
让我们再算一次:
nums = [1, -2, 3, 4, -1, 2, 1, -5, 4]
子数组 [3, 4, -1, 2, 1, -5, 4] 长度 7,和 = 3+4-1+2+1-5+4 = 8
不对。
也许答案是 9?但示例说 10。
我们手动算一下所有至少长度为 3 的子数组:
[1,-2,3] = 2
[-2,3,4] = 5
[3,4,-1] = 6
[4,-1,2] = 5
[-1,2,1] = 2
[2,1,-5] = -2
[1,-5,4] = 0
[1,-2,3,4] = 6
[-2,3,4,-1] = 4
[3,4,-1,2] = 8
[4,-1,2,1] = 6
[-1,2,1,-5] = -3
[2,1,-5,4] = 2
[1,-2,3,4,-1] = 5
[-2,3,4,-1,2] = 6
[3,4,-1,2,1] = 9
[4,-1,2,1,-5] = 1
[-1,2,1,-5,4] = 1
更长的不可能超过 9。
所以示例的输出 10 是错的,应该是 9。我们继续用正确的思路。
修正递推公式
我们的递推公式 dp[j] = max( sum[j-k+1 .. j], dp[j-1] + nums[j] ) 正确吗?
考虑 j=5 时,dp[5] 应该是多少?
至少长度为 3 的子数组以索引 5 结尾:
可能子数组:
[3,4,-1,2] 和 8
[4,-1,2] 和 5
[-1,2] 长度不够 3,不考虑
[2] 不考虑
所以最大和是 8。
我们的公式:
sum[3..5] = nums[3]+nums[4]+nums[5] = 4 + (-1) + 2 = 5
dp[4] + nums[5] = 6 + 2 = 8
max(5,8)=8 ✅
所以公式是对的,但之前我计算 dp[4] 时错了。我们重新正确计算一遍。
重新计算示例
nums = [1, -2, 3, 4, -1, 2, 1, -5, 4], k=3
prefix = [0, 1, -1, 2, 6, 5, 7, 8, 3, 7]
j=2: dp[2] = sum[0..2] = prefix[3]-prefix[0]=2-0=2
j=3: sum[1..3]= prefix[4]-prefix[1] = 6-1=5
dp[2]+nums[3] = 2+4=6
dp[3]=max(5,6)=6
j=4: sum[2..4]= prefix[5]-prefix[2] = 5-(-1)=6
dp[3]+nums[4] = 6+(-1)=5
dp[4]=max(6,5)=6
j=5: sum[3..5]= prefix[6]-prefix[3] = 7-2=5
dp[4]+nums[5] = 6+2=8
dp[5]=max(5,8)=8
j=6: sum[4..6]= prefix[7]-prefix[4] = 8-6=2
dp[5]+nums[6] = 8+1=9
dp[6]=max(2,9)=9
j=7: sum[5..7]= prefix[8]-prefix[5] = 3-5=-2
dp[6]+nums[7] = 9+(-5)=4
dp[7]=max(-2,4)=4
j=8: sum[6..8]= prefix[9]-prefix[6] = 7-7=0
dp[7]+nums[8] = 4+4=8
dp[8]=max(0,8)=8
dp数组: [2,6,6,8,9,4,8]
最大值是 9。
所以答案是 9。
最终算法步骤
- 如果数组长度 n < k,无法满足至少 k 个元素,但题目一般保证 k ≤ n。
- 计算前缀和数组 prefix,长度为 n+1。
- 初始化 dp[k-1] = prefix[k] - prefix[0]。
- 初始化答案 ans = dp[k-1]。
- 对于 j 从 k 到 n-1:
- currentKsum = prefix[j+1] - prefix[j-k+1]
- extend = dp[j-1] + nums[j]
- dp[j] = max(currentKsum, extend)
- 更新 ans = max(ans, dp[j])
- 返回 ans。
复杂度分析
- 时间复杂度:O(n),我们只遍历了一次数组。
- 空间复杂度:O(n),用于存储前缀和和 dp 数组。可以优化到 O(1) 只存储前一个 dp 值,因为递推只用到 dp[j-1]。
空间优化
我们只需要维护一个变量 prev_dp 表示 dp[j-1],以及一个变量 ans 记录最大值。
算法(空间优化版):
1. 计算前缀和数组 prefix(长度 n+1)。
2. 初始化 prev_dp = prefix[k] - prefix[0]。
3. ans = prev_dp。
4. 对于 j 从 k 到 n-1:
currentKsum = prefix[j+1] - prefix[j-k+1]
extend = prev_dp + nums[j]
cur_dp = max(currentKsum, extend)
ans = max(ans, cur_dp)
prev_dp = cur_dp
5. 返回 ans。
边界情况
- 所有元素都是负数:我们必须选择一个长度至少为 k 的子数组,所以答案是“绝对值最小的负数之和”,即最大的那个长度为 k 的子数组和,或者更大的扩展如果能使和更大(实际上对于全负数,扩展只会让和更小,所以最优就是长度为 k 的最大和子数组)。
- k = 1:退化为普通的 Kadane 算法求最大子数组和。
- k = n:只能选择整个数组,答案就是整个数组的和。
总结
这道题通过动态规划巧妙地将“至少包含k个元素”的条件转化为两种情况的比较:要么取长度为k的窗口,要么从前一个至少长度为k的最优解扩展而来。这种思路将复杂度从 O(n²) 降到了 O(n),是处理此类长度限制子数组和问题的典型方法。