每周一算法：二维前缀和

二维前缀和

对一个序列预处理得到前缀和数组，可以在$O(1)$的时间复杂度计算序列中任意区间的元素之和，这是前缀和算法的作用。而二维前缀和是用来优化处理子矩阵的和。

例如，对于矩阵$A=\left[\begin{array}{cccc}1& 4& 5& 2\\ 9& 5& 2& 1\\ 6& 9& 8& 3\\ 4& 2& 1& 6\end{array}\right]$中的一个子矩阵$\left[\begin{array}{cc}5& 2\\ 9& 8\end{array}\right]$的所有元素之和为$24$。

使用二维前缀和可以在$O(1)$的时间复杂度计算出该子矩阵的和。

算法思想

构造前缀和数组

首先需要在原矩阵的基础上预处理一个二维前缀和数组$s[][]$，其中$s[i][j]$表示从第$1$行第$1$列开始到第$i$行第$j$列所有数的和。

例如，对于矩阵$A=\left[\begin{array}{cccc}1& 4& 5& 2\\ 9& 5& 2& 1\\ 6& 9& 8& 3\\ 4& 2& 1& 6\end{array}\right]$：

• $s[1][1]$表示矩阵第$1$行第$1$列所有数的和，为$1$；
• $s[1][2]$表示矩阵第$1$行第$1$列到第$1$行第$2$列所有数的和，为$1+4=5$；
• …
• $s[3][3]$该如何计算呢？

分析下图可以发现，$s[3][3]$由三部分组成

• 绿色部分$s[2][3]$
• 红色部分$s[3][2]$
• 紫色部分$A[3][3]$。

由于$s[2][2]$（黄色部分）被包含了两次，因此在计算$s[3][3]$时，需要去掉一个$s[2][2]$。因此，$s[3][3]=s[2][3]+s[3][2]-s[2][2]+A[3][3]$。

那么，从第$1$行第$1$列开始到第$i$行第$j$列所有数的和$s[i][j]=s[i-1][j]+s[i][j-1]-s[i-1][j-1]+A[i][j]$；

计算子矩阵的和

有了前缀和数组，如何计算矩阵$A=\left[\begin{array}{cccc}1& 4& 5& 2\\ 9& 5& 2& 1\\ 6& 9& 8& 3\\ 4& 2& 1& 6\end{array}\right]$中的一个子矩阵$\left[\begin{array}{cc}5& 2\\ 9& 8\end{array}\right]$的所有元素之和？

通过下图可以发现，从$(2,2)$到$(3,3)$的子矩阵之和可以由下面几个部分的和相减得到：

• 天蓝色部分$s[3][3]$
• 绿色部分 $s[1][3]$
• 红色部分$s[3][1]$

由于黄色部分$s[1][1]$被减了两次，因此计算结果时需要补加一次$s[1][1]$。那么从$(2,2)$开始到$(3,3)$结束的子矩阵之和$=s[3][3]-s[1][3]-s[3][1]+s[1][1]$。

更一般地，从$(x1,y1)$开始到$(x2,y2)$结束的子矩阵之和$=s[x2][y2]-s[x1-1][y2]-s[x2][y1-1]+s[x1-1][y1-1]$。

算法实现

//构造前缀和数组
for(int i = 1; i <= n; i ++)
	for(int j = 1; j <= m; j ++)
	 	s[i][j] = s[i -1][j] + s[i][j - 1] - s[i - 1][j - 1] + a[i][j];
//输出(x1,y1)到(x2,y2)子矩阵的和
cout << s[x2][y2] - s[x1 - 1][y2] - s[x2][y1 - 1] + s[x1 - 1][y1 - 1] << endl;	 

时间复杂度

二维前缀和算法的时间复杂度分为两部分：

• 构造二维前缀和数组：$O(n\times m)$
• 求子矩阵的和：$O(1)$

空间复杂度

需要构造和原矩阵大小一样的前缀和数组，因此空间复杂度为$O(n\times m)$。

真题演练

洛谷P2004 领地选择

题目描述

作为在虚拟世界里统帅千军万马的领袖，小 Z 认为天时、地利、人和三者是缺一不可的，所以，谨慎地选择首都的位置对于小 Z 来说是非常重要的。

首都被认为是一个占地 $C\times C$ 的正方形。小 Z 希望你寻找到一个合适的位置，使得首都所占领的位置的土地价值和最高。

输入格式

第一行三个整数 $N,M,C$，表示地图的宽和长以及首都的边长。

接下来 $N$ 行每行 $M$ 个整数，表示了地图上每个地块的价值。价值可能为负数。

输出格式

一行两个整数 $X,Y$，表示首都左上角的坐标。

样例输入

3 4 2
1 2 3 1
-1 9 0 2
2 0 1 1

样例输出

1 2

提示

对于 $60\%$ 的数据，$N,M\le 50$。

对于 $90\%$ 的数据，$N,M\le 300$。

对于 $100\%$ 的数据，$1\le N,M\le 10^3$，$1\le C\le \min (N,M)$。

解题思路

根据题目描述，在$N\times M$的矩阵中求大小为 $C\times C$ 的子矩阵的土地价值和的最大值，因此可以使用二维前缀和优化。

实现代码

#include <iostream>
using namespace std;
const int N = 1010;
long long s[N][N];
int n, m, c, x;
int main()
{
    cin >> n >> m >> c;
    for(int i = 1; i <= n; i ++)
        for(int j = 1; j <= m; j ++)
        {
            scanf("%d", &x);
            s[i][j] = s[i - 1][j] + s[i][j - 1] - s[i - 1][j - 1] + x; 
        }
    long long ans = -1e10;
    int a, b;
    //枚举子矩阵的开始位置，注意子矩阵不能越界
    for(int x1 = 1; x1 + c - 1 <= n; x1 ++)
        for(int y1 = 1; y1 + c - 1 <= n; y1 ++)
        {
            int x2 = x1 + c - 1, y2 = y1 + c - 1;
            int t = s[x2][y2] - s[x1 - 1][y2] - s[x2][y1 - 1] + s[x1 - 1][y1 - 1];
            if(t > ans)
            {
                ans = t, a = x1, b = y1;
            }
        }
    cout << a << ' ' << b << endl;
}

总结

• 二位前缀和可以优化求子矩阵的和
• 构造二维前缀和数组：s[i][j] = s[i -1][j] + s[i][j - 1] - s[i - 1][j - 1] + a[i][j];
• 使用前缀和数组求子矩阵的和：s[x2][y2] - s[x1 - 1][y2] - s[x2][y1 - 1] + s[x1 - 1][y1 - 1]