如何以更优时间复杂度求解M次前缀和计算问题

O(NM)的动态规划之所以超时，本质是没有利用到多次前缀和的组合数性质，下面给出的解法和M的大小完全无关，哪怕M取值到1e9也可以正常计算。

核心结论 #

经过M次前缀和操作后，原数组位置j的元素，对最终结果数组位置i（i >= j）的贡献权重为组合数C(M + i - j - 1, i - j)，即：
res[i] = sum_{j=0}^i a[j] * C(M + i - j - 1, i - j)

用题目给出的两个样例代入计算，结果和样例输出完全匹配。

系数递推方法 #

不需要提前预处理大范围阶乘，权重可以直接线性递推算出：记间隔d = i-j，对应权重c[d] = C(M + d -1, d)，递推规则为：

• 初始值c[0] = 1（间隔为0时，元素对自身位置的贡献永远是1）
• 对d从1到N-1：c[d] = c[d-1] * (M + d - 1) // d
  递推得到的所有c[d]都是整数，直接做整数除法即可，没有精度问题。

代码实现 #

以下是Python版本的可运行代码，完全适配题目给出的两个样例：

def m_times_prefix_sum(arr, M):
    n = len(arr)
    if M == 0 or n == 0:
        return arr.copy()
    # 线性递推计算权重数组
    c = [0] * n
    c[0] = 1
    for d in range(1, n):
        c[d] = c[d-1] * (M + d - 1) // d
    # 计算最终结果
    res = [0] * n
    for i in range(n):
        s = 0
        for d in range(i+1):
            s += arr[i - d] * c[d]
        res[i] = s
    return res

# 样例1测试
print(m_times_prefix_sum([1,2,3], 4))  # 输出 [1, 6, 21]
# 样例2测试
print(m_times_prefix_sum([1,2,3,4,5], 3)) # 输出 [1, 5, 15, 35, 70]

复杂度说明 #

• 上述基础实现的时间复杂度为O(N²)，空间复杂度为O(N)，运算效率和M的取值完全无关，哪怕M取1e9也可以在相同时间内算出结果，当N规模在1e4以内时运行速度远快于O(NM)的动态规划方法，不会出现超时问题。
• 如果N规模达到1e5级别，可以将权重数组反转后和原数组做FFT/NTT快速卷积，取前N项结果即可，时间复杂度可以进一步优化到O(N log N)。

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