【算法】前缀和算法详解

1. 引入

我们先来看这样一个问题：给定一个数组 a = [1, 2, 3, 4, 5, 6]，如何求解数组中所有元素的和？

这太简单了，我们只需要枚举每一个元素，然后依次相加即可。

int sum = 0;
for(int i = 0; i < n; i++){
	sum += a[i];
}

那如果我只想要数组中其中一部分的和呢？比如：[3, 4, 5] 这一部分。

那也不难，我们只需要枚举从区间开始到结束的所有元素的和即可。

int sum = 0;
for(int i = L; i <= R; i++){
	sum += a[i];
}

那如果如果有成千上万次这样的询问区间的次数，我们还是这样处理吗？

while(m--){  // m表示询问次数
    int sum = 0;
    for (int i = L; i <= R; i++) {
		sum += a[i];
	}
}

如果是这样的话，时间复杂度就变成了计算区间长度乘以询问次数 $O(n\ast m)$，并不是很高效，于是我们便可以利用一种询问区间和的高效算法——前缀和。

2. 一维前缀和

2.1 算法介绍

上面的做法中，sum 记录的是 a[0]+a[1]+...+a[i]，也就是前 i 个数的和，所以我们可以定义一个数组prefix[]，将每次计算出来的这个 sum 储存在这个数组 prefix[] 中，prefix[i] 就代表 a[] 数组中前 i 个元素的和。

int sum = 0;
for(int i = 0;i < n; i++){
    sum += a[i];
    prefix[i] = sum;
}

这个时候如果我们再想要去求区间的和，比如区间 [5, 10] 时，就只需要用 prefix[10]-prefix[4]即可。
因为 prefix[10] 是原数组中下标从 0 到 10 的数之和，prefix[5] 是下标 0 到 4 的数之和，相减之后自然就是下标 5 到 10 的和。

这也就意味着我们只需要先用 $O(n)$ 的时间去创建一个前缀和数组 prefix[]，然后只需要用 $O(1)$ 的时间去求得任意子区间的和。
综上所述，可以得到一个公式：sum[L, R] = prefix[R] - prefix[L - 1]

2.2 优化

上面的代码还可以进一步优化，可以像下面这样写，就可以避免引入一个 sum 变量，使得代码更简洁。

for(int i = 1; i < n; i++){
    prefix[i] = prefix[i - 1] + a[i];
}

注意这样写的话那么下标要从 1 开始写，因为这种写法你会发现有一个 [i - 1]，如果 i == 0 的话就会出现越界的问题，那么这个时候 prefix[0] 该怎么处理？为了解决这个问题我们可以事先让 prefix[0] = 0，相当于初始化添加了一个虚拟的节点（辅助节点）。

2.3【模板】前缀和

题目链接——【模板】前缀和

#include <stdio.h>

int main() {
    int n, q;  // 第一行要输入的整数
    scanf("%d %d", &n, &q);
    // 可能有些编译器不支持使用变量作为数组大小，你也可以定义一个宏，不方便的话直接去用C++写吧...
    int a[n + 1];  // 要输入的数组， 开 n+1 个是因为第一个位置放0
    a[0] = 0;
    long long prefix[n + 1];  // 前缀和数组同样开 n+1 个
    prefix[0] = 0;
    for(int i = 1; i < n + 1; ++i){  // 注意下标从 1 开始
        scanf("%d", a + i);
        prefix[i] = prefix[i - 1] + a[i];  // 前缀和
    }
    int l, r;  // 区间的左右边界
    long long sum;  // 区间和
    while(q--){  // 访问 q 次
        scanf("%d %d", &l, &r);
        sum = prefix[r] - prefix[l - 1];  // 直接用公式求区间和
        printf("%lld\n", sum);
    }
    return 0;
}

2.4【扩展】前缀积，前缀异或

前缀和这种思想不仅仅可以用于求和，也可以用于求其他更多的运算，比如乘法、异或等。因为它们都满足结合律，且它们是可逆运算。

• 前缀积

之前我们计算区间和的公式是这样的：
$$\mathrm{sum}[L,R]=\mathrm{prefix}[R]-\mathrm{prefix}[L-1]$$
而现在对于一个区间的乘法而言，我们可以做一个小小的变形即可：
$$\mathrm{mul}[L,R]=\mathrm{prefix}[R]/\mathrm{prefix}[L-1]$$
即用 R 的前缀积除以 L - 1 的前缀积，但前提是前缀积不能溢出。由于这样的前缀积的结果非常大，为了避免高精度的计算，通常我们可以对结果进行取模，如下：
$$\mathrm{prefix}[i]=(\mathrm{prefix}[i-1]\ast a[i])\mathrm{mod}p$$
但如果是这样的话，它的逆运算就不像加法对应减法、乘法对应除法这么简单。它的逆运算需要了解 “乘法逆元” 的相关知识，这里就不展开了，有兴趣的话可以看看我写的杂文里有关数论的内容。

• 前缀异或

由异或的自反性可得，异或的逆运算其实就是本身，因为 a ^ b ^ b = a。所以有：
$$\mathrm{xor}[L,R]=\mathrm{prefix}[R]\oplus \mathrm{prefix}[L-1]$$
有了这个前缀异或，加入说我们在一个前缀异或数组中存在两个相同的数，下标为 $l$ 和 $r$，那么说明在原数组区间 $[l+1,r]$ 中的这些数的异或和为 0。如果前缀异或数组中有多个数相同，那么任选两个数它们对应到原数组中的区间，这个区间内的数异或和为 0。通过这个性质，我们就可以通过排列组合快速求得数组中有多少段区间的异或和为 0。注意：数组最开头的虚拟节点 0 也要加上。

3. 二维前缀和

3.1 算法介绍

和一维前缀和类似，也是开辟一个前缀和数组，只不过这次是二维的，相当于一个矩阵。我们对行和列都使用一维前缀和的算法，便可以得到一个二维前缀和矩阵 prefix[][]。不难发现，这个 prefix 矩阵的元素 prefix[i][j] 其实就是从左上角开始到该位置所框出来的矩形内所有数据的和。

当然我们不可能一个一个枚举原数组中的每一个元素，根据前缀和的原理，就有：prefix[i][j] = a[i][j] + prefix[i - 1][j] + prefix[i][j - 1] - prefix [i - 1][j - 1]。

那么如果让你求原矩阵中任意一个子矩阵的和呢？比如给你一个子矩阵左上角的坐标位置以及右下角的坐标位置，如下：

其实要理解起来的话也不难，本质上就是一个容斥原理。那么根据这个原理，假如说给你左上角的坐标为（x1, y1），右下角的坐标为（x2, y2），那么就有：子矩阵之和 = prefix[x2][y2] - prefix[x2][y1 - 1] - prefix[x1 - 1][y2] + prefix[x1 - 1][x2 - 1]

3.2【模板】二维前缀和

题目链接——【模板】二维前缀和

#include<stdio.h>
#define N 1001

int main() {
    int n, m, q;  // 第一行的三个数 
    scanf("%d %d %d", &n, &m, &q);
    // 这里我们用 temp 就表示 a[i][j] 的值，不需要再开一个数组a
    long long temp, prefix[N][N] = {0};  // 前缀和数组初始化为0
    for(int i = 1; i <= n; i++){  // 注意下标从 1 开始
        for(int j = 1; j <= m; j++){
            scanf("%lld", &temp);  // 套公式
            prefix[i][j] = temp + prefix[i - 1][j] + prefix[i][j - 1] - prefix[i - 1][j - 1];
        }
    }
    int x1, y1, x2, y2; 
    long long res;
    while(q--){
        scanf("%d %d %d %d", &x1, &y1, &x2, &y2);
        res = prefix[x2][y2] - prefix[x1 - 1][y2] - prefix[x2][y1 - 1] + prefix[x1 - 1][y1 - 1];
        printf("%lld\n", res);
    }
    return 0;
}

4. OJ 实战

4.1 寻找数组的中心下标

题目链接——【寻找数组的中心下标】

这道题目不仅会用到原数组从前往后的元素和，也会用到从后往前的所有元素和，因此我们可以定义两个数组，一个前缀数组，一个后缀数组。分别用来记录两部分的区间和。

在这里会发现和最开始讲的一维前缀和略有不同，这里的 prefix[i] = prefix[i - 1] + nums[i - 1] 而不是 nums[i]，后缀和同理。这是因为这道题原数组是已经给你了的，下标是从 0 开始的，所以加 nums 数组中的数时要往前移一个位置。

还有一个地方需要说明的是，你会发现前后缀这两个数组开辟和 nums 同样的大小的话，这两个数组是记录不到 nums 的所有元素之和的。但是这道题恰好也算不到所有元素的和，最多只能算 numsSize - 1个数之和，所以开同样大小刚好符合题意。

int pivotIndex(int* nums, int numsSize) {
    int prefix[numsSize];  // 定义前缀和数组
    prefix[0] = 0;  // 初始化第一个位置为0
    int suffix[numsSize];  // 后缀数组
    suffix[numsSize - 1] = 0;  // 初始化最后一个位置为0
    
    for(int i = 1; i < numsSize; i++){  // 开始求前缀和
        prefix[i] = prefix[i-1]+nums[i-1];  // 注意是nums[i-1]
    }
    
    for(int i = numsSize - 2; i >= 0; i--){  // 求后缀和
        suffix[i] = suffix[i+1]+nums[i+1];   
    }
    
    for(int i = 0; i < numsSize; i++){
        if(prefix[i] == suffix[i]){
            return i;
        }
    }
    
    return -1;  // 出循环了说明没找到，返回-1
}

4.2 除自身以外数组的乘积

4.2.1 常规

题目链接——【除自身以外数组的乘积】

和上一道题非常类似，只不过这道题求的是积，由于题目不让使用除法，那么很容易想到前缀和的扩展——前缀积。

需要注意的是这里两个数组开头和结尾的辅助节点需要初始化为 1，要不然后面的就全是 0 了。

/**
 * Note: The returned array must be malloced, assume caller calls free().
 */
int* productExceptSelf(int* nums, int numsSize, int* returnSize) {
    int prefix[numsSize];  // 前缀积数组
    prefix[0] = 1;  // 因为这里是乘法运算，所以要初始化为1
    int suffix[numsSize];  // 后缀积数组
    suffix[numsSize - 1] = 1;  // 最后一个位置初始化为1

    for(int i = 1; i < numsSize; i++){  // 求前缀积
        prefix[i] = prefix[i-1] * nums[i-1];
    }
    
    for(int i = numsSize - 2; i >= 0; i--){  // 求后缀积
        suffix[i] = suffix[i+1] * nums[i+1];
    }

    int* ret = (int*)malloc(numsSize * sizeof(int));  // 要求返回的数组
    *returnSize = numsSize;  // 该数组的大小
    for(int i = 0; i < numsSize; i++){
        ret[i] = prefix[i] * suffix[i];
    }

    return ret;
}

4.2.2 优化

上述解法的空间复杂度为 $O(n)$，我们还可以将其优化至 $O(1)$。

我们可以指开辟一个后缀积数组计算原数组的后缀积，然后再用一个循环从前往后遍历原数组，一边计算前缀积一遍把这个前缀积乘到后缀积数组中，这样这个后缀积数组就变成了题目要求返回的数组了。

而由于返回的数组不算作额外的空间，所以空间复杂度为 $O(1)$。

int* productExceptSelf(int* nums, int numsSize, int* returnSize) {
    int* suffix = (int*)malloc(numsSize * sizeof(int));  // 后缀积数组
    suffix[numsSize - 1] = 1;
    for(int i = numsSize - 2; i >= 0; i--){
        suffix[i] = suffix[i+1] * nums[i+1];
    }

    int pre = 1;
    for(int i = 0; i < numsSize; i++){
        // 此时 pre 为 nums[0] 到 nums[i-1] 的乘积，直接乘到 suf[i] 中
        suffix[i] *= pre;
        pre *= nums[i];
    }

    *returnSize = numsSize;
    return suffix;
}

4.3 矩阵区域和

题目链接——【矩阵区域和】

先来理解一下题意：

所以说这道题的本质就是在求子矩阵的和，可以用到二维前缀和的技巧。

• 思路

我们先预处理出一个前缀和矩阵然后根据子矩阵左上角的坐标（x1, y1）和右下角的坐标 （x2, y2）来计算子矩阵的和从而对应到要返回的 ans 矩阵中。

• 细节

1. 之前我们所写的二维前缀和模板是 prefix[i][j] = a[i][j] + prefix[i - 1][j] + prefix[i][j - 1] - prefix [i - 1][j - 1]。但是注意！这里的二维数组 mat 的下标是从 0 开始的，和前面的一维前缀和的题目类似，所以这里的公式里不能加 a[i][j]，而是 a[i - 1][j - 1]。
   所以有 prefix[i][j] = a[i - 1][j - 1] + prefix[i - 1][j] + prefix[i][j - 1] - prefix [i - 1][j - 1]。

2. 如何求得子矩阵的坐标（x1, y1）和右下角的坐标 （x2, y2）？
   很简单，根据 k 的值，在 mat 矩阵中的某个元素 mat[i][j] 所向外扩展出来的矩阵的左上角坐标就是 [i - k]|[j - k]，那么同理右下角的坐标就是 [i + k][j + k]。

   但是这里有个问题就是越界，比如上面的图中明显红色方框就框出界了，这个时候我们只需要限制一下 x1,x2, y1, y2 的大小即可，比如 x1 在前缀和矩阵中最小是 1，那么就只需要让 x1 = fmax(1, x1) 即可。加入列数为 col，那么 y2 最大是 col，只需让 y2 = fmin(col, y1) 即可。

这里解释一下这道题的函数中的各个参数的意思：

1. 首先返回的是一个矩阵，C 语言中可以用二级指针来表示。
2. int** mat：二级指针，表示原矩阵，即示例给你的矩阵。
3. int matSize：由于矩阵实际上是用一维数组来表示的，而这个一维数组中的元素又是多个一维数组，这样才表示的是一个二维数组，即矩阵。这个参数就是说这个外层的一维数组有多少个元素，对应到矩阵中就是矩阵的行数。
4. int* matColSize：一级指针，表示一个数组，存储的是这个矩阵的每一行的列数，由于这里是矩阵，肯定是矩形的，所以每一行的列数都相等，可以用 matColsize[0] 来获取这个矩阵的列数。
5. int k：表示向四周延伸的距离。
6. int* returnSize：一个指向整数的指针，通过这个参数来返回答案矩阵的行数。
7. int** returnColumnSizes：这个参数用来返回一个数组，数组中每一个元素代表返回的矩阵每一行的列数。

/**
 * Return an array of arrays of size *returnSize.
 * The sizes of the arrays are returned as *returnColumnSizes array.
 * Note: Both returned array and *columnSizes array must be malloced, assume caller calls free().
 */

// 定义二维前缀和数组 
int prefix[101][101];
// 这个函数用于求子矩阵的和
int Sum(int x1, int y1, int x2, int y2, int row, int col){
    x1 = fmax(1, x1);  // 四个语句处理边界情况
    y1 = fmax(1, y1);
    x2 = fmin(row, x2);
    y2 = fmin(col, y2);
    // 子矩阵的和（公式）
    return prefix[x2][y2] - prefix[x1 - 1][y2] - prefix[x2][y1 - 1] + prefix[x1 - 1][y1 - 1];
}

int** matrixBlockSum(int** mat, int matSize, int* matColSize, int k, int* returnSize, int** returnColumnSizes) {
    // 定义矩阵的行和列数
    int row = matSize, col = matColSize[0];
    // 开始预处理前缀和矩阵
    for(int i = 1; i <= row; i++){
        for(int j = 1; j <= col; j++){
            // 注意是 mat[i - 1][j - 1]，因为 mat 下标是从 0 开始
            prefix[i][j] = prefix[i - 1][j] + prefix[i][j - 1] - prefix[i - 1][j - 1] + mat[i - 1][j - 1];
        }
    }
	// 根据题意，返回的数组要自己 malloc，且不用自己 free
    int** ans = (int**)malloc(sizeof(int*) * row); 
    *returnSize = row; 
    *returnColumnSizes = matColSize;
    
 	// 二维数组的 malloc 需要对数组内的每一个一维数组进行开辟空间的操作
    for(int i = 0; i < row; i++){
        ans[i] = (int*)malloc(sizeof(int) * col);
    }
	
    // 开始使用前缀和数组
    for(int i = 1; i <= row; i++){
        for(int j = 1; j <= col; j++){
            ans[i - 1][j - 1] = Sum(i - k, j - k, i + k, j + k, row, col);
        }
    }
    return ans;
}