如何优化k-最大子数组和问题的O(nk)时间复杂度代码？

问题描述 #

给定长度为n的整数数组A，k-最大子数组和问题要求找出最多k个不重叠的连续子数组，使其总和最大；若数组全为负数，总和为0。例如数组A=[-1,2,-1,2,-1,2,2]时，两个最优子数组为[2,-1,2]和[2,2]，总和为7。

现有O(nk)复杂度实现及问题 #

我基于Kadane算法扩展实现了时间复杂度为O(nk)的代码，如下：

import numpy as np

def solve_SO(test_seq, k=2):
    """
    Computes the k max subarray sum
    """
    num_intervals = k * 2 + 1
    best = np.zeros(num_intervals, dtype=int)
    for seq_idx, val in enumerate(test_seq):
        # 将当前值加到所有"包含当前元素"的区间最优得分中
        for interval_idx in range(1, num_intervals, 2):
            best[interval_idx] += val
        # 遍历所有区间，若当前区间得分不如前一个，则更新为前一个的得分，保证best数组单调递增
        for interval_idx in range(1, num_intervals):
            if best[interval_idx] < best[interval_idx - 1]:
                best[interval_idx] = best[interval_idx - 1]
    return best[num_intervals - 1]

测试示例 #

solve_SO([-1, 2, -1, 2, -1], 2)
# 输出：4

solve_SO([-1, 2, -1, 2, -1], 1)
# 输出：3

但该代码在k较大时效率较低，比如k=3时会生成长度为7的辅助数组，且遍历输入数组的每个元素时都要遍历整个辅助数组。

优化方案 #

1. 压缩状态数组，减少内层循环次数 #

原代码维护了长度为2k+1的状态数组，实际上可以简化状态定义，只维护两组核心状态：

• dp[j]：前i个元素中选j个不重叠子数组的最大和（当前不处于子数组中）
• curr[j]：前i个元素中选j个不重叠子数组的最大和（当前处于子数组中）

通过倒序更新这两个数组（避免覆盖未使用的状态），可以将内层循环次数从2k减少到k，同时保持O(k)的空间复杂度，实际运行效率更高：

def max_k_subarray_sum(arr, k):
    n = len(arr)
    # dp[j]：选j个子数组的最大和（当前不在子数组）
    dp = [0] * (k + 1)
    # curr[j]：选j个子数组且当前处于子数组中的最大和
    curr = [-float('inf')] * (k + 1)
    
    for num in arr:
        # 倒序更新，防止覆盖还未使用的上一轮状态
        for j in range(k, 0, -1):
            # 要么继续当前子数组，要么从"选j-1个的最大和"开始新子数组
            curr[j] = max(curr[j] + num, dp[j-1] + num)
            # 要么保持之前的最优解，要么结束当前子数组（取curr[j]的最大值）
            dp[j] = max(dp[j], curr[j])
    
    # 取最多k个的最大值，全负时返回0
    return max(max(dp), 0)

2. 大k值场景的特殊处理 #

当k >= len(arr)//2时，最优解等价于取数组中所有正数的和（因为每个正数都可以单独作为一个子数组，且k足够容纳所有正数）。此时无需执行O(nk)的循环，直接计算正数和即可，时间复杂度降为O(n)：

def max_k_subarray_sum(arr, k):
    # 大k值直接取所有正数和
    positive_sum = sum(x for x in arr if x > 0)
    if k >= len(arr) // 2:
        return max(positive_sum, 0)
    
    n = len(arr)
    dp = [0] * (k + 1)
    curr = [-float('inf')] * (k + 1)
    
    for num in arr:
        for j in range(k, 0, -1):
            curr[j] = max(curr[j] + num, dp[j-1] + num)
            dp[j] = max(dp[j], curr[j])
    
    return max(max(dp), 0)

3. 移除numpy依赖提升小k场景效率 #

原代码使用numpy数组实现状态维护，但对于小k值的场景，普通Python列表的操作开销更低，且无需引入额外依赖，进一步提升运行速度。

内容的提问来源于stack exchange，提问作者Simd