如何在O(N+M)时间复杂度下完成N*M元素的双重求和？

近期遇到一道算法题，题目定义了一个操作两个数组的函数（函数定义见下图）：

我通过两层循环实现了该函数，代码如下：

int f(int n, int m, int *a, int *b)
{
    int x = 0;
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < m; j++) {
            x += a[i] * b[j] * (i + j + 2); // (... + 2) 补偿0-based索引
        }
    }    

    return x % 998244353; // 题目指定的常数
}

但题目要求实现线性时间复杂度（数组A和B大小的线性级别）的解决方案，我思考许久仍不知如何将两层循环简化，求解决建议。

核心是通过代数拆分将双重循环的累加式拆解为多个独立的单数组求和，把O(n*m)复杂度降到O(n+m)。

1. 拆分累加项 #

先把原式中的(i + j + 2)拆分为(i+1)+(j+1)（对应题目中的1-based索引，代码用0-based所以加1还原），代入原累加式后拆分：

x = Σ(i=0到n-1) Σ(j=0到m-1) a[i]*b[j]*(i+1) + Σ(i=0到n-1) Σ(j=0到m-1) a[i]*b[j]*(j+1)

由于两个求和项中的变量可以分离，进一步转化为：

x = [Σa[i]*(i+1)] * [Σb[j]] + [Σa[i]] * [Σb[j]*(j+1)]

2. 计算四个线性求和值 #

只需要遍历两次数组，计算四个独立的累加结果：

• sum_a：数组a所有元素之和
• sum_a_weighted：数组a每个元素乘以(i+1)的和（i为0-based索引）
• sum_b：数组b所有元素之和
• sum_b_weighted：数组b每个元素乘以(j+1)的和（j为0-based索引）

3. 线性时间代码实现 #

用long long避免整数溢出，每一步取模保证数值在题目指定的998244353范围内：

const int MOD = 998244353;

int f(int n, int m, int *a, int *b)
{
    long long sum_a = 0, sum_a_weighted = 0;
    for (int i = 0; i < n; i++) {
        sum_a = (sum_a + a[i]) % MOD;
        sum_a_weighted = (sum_a_weighted + (long long)a[i] * (i + 1)) % MOD;
    }

    long long sum_b = 0, sum_b_weighted = 0;
    for (int j = 0; j < m; j++) {
        sum_b = (sum_b + b[j]) % MOD;
        sum_b_weighted = (sum_b_weighted + (long long)b[j] * (j + 1)) % MOD;
    }

    long long x = (sum_a_weighted * sum_b % MOD + sum_a * sum_b_weighted % MOD) % MOD;
    return (int)x;
}

复杂度验证 #

遍历数组a耗时O(n)，遍历数组b耗时O(m)，总时间复杂度为O(n+m)，完全满足题目要求的线性级别。

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