前缀和算法实战:连续数组与矩阵区域和
前缀和是处理区间求和问题的高效技巧,广泛应用于一维数组及二维矩阵的场景。本文将通过两道经典题目,深入讲解如何利用哈希表优化一维前缀和查询,以及如何构建二维前缀和矩阵实现 O(1) 的区域求和。
031 连续数组
题目描述:
给定一个二进制数组 nums,找到含有相同数量的 0 和 1 的最长连续子数组,并返回该子数组的长度。
思路分析
暴力解法需要枚举所有子数组并统计 0 和 1 的数量,时间复杂度为 O(n²),在数据量较大时容易超时。我们需要寻找更优的解法。
核心转化在于将问题简化:如果我们将所有的 0 视为 -1,那么原问题就转化为'寻找和为 0 的最长子数组'。这与经典的「和为 K 的子数组」问题思路一致。
设 sum[i] 表示从索引 0 到 i 的前缀和。若存在两个位置 i 和 j(假设 i < j),使得 sum[j] == sum[i],则说明区间 (i, j] 内的元素之和为 0。因此,为了找到最长的满足条件的子数组,我们需要记录每个前缀和第一次出现的位置。
我们不需要显式地维护一个前缀和数组,只需使用哈希表存储 {前缀和值:首次出现的索引}。初始化时,哈希表中存入 {0: -1},代表前缀和为 0 出现在索引 -1 处(即空数组状态),这样当遍历到某个位置前缀和恰好为 0 时,可以直接计算长度。
代码实现
C++ 实现
class Solution {
public:
int findMaxLength(vector<int>& nums) {
unordered_map<int, int> hash;
// 初始化:前缀和为 0 时,对应索引为 -1
hash[0] = -1;
int sum = 0, ret = 0;
for (int i = 0; i < nums.size(); ++i) {
// 0 记为 -1,1 记为 1
sum += (nums[i] == 0 ? -1 : 1);
if (hash.count(sum)) {
// 如果当前前缀和之前出现过,更新最大长度
ret = max(ret, i - hash[sum]);
} else {
// 否则记录该前缀和首次出现的位置
hash[sum] = i;
}
}
return ret;
}
};
复杂度分析:
- 时间复杂度:O(n),仅需遍历一次数组。
- 空间复杂度:O(n),哈希表在最坏情况下需存储 n 个不同的前缀和。
Java 实现
class Solution {
public int findMaxLength(int[] nums) {
Map<Integer, Integer> hash = new HashMap<>();
// 默认存在一个前缀和为 0 的情况,对应索引 -1
hash.put(0, -1);
int sum = 0, ret = 0;
for (int i = 0; i < nums.length; i++) {
sum += (nums[i] == 0 ? -1 : 1);
if (hash.containsKey(sum)) {
ret = Math.max(ret, i - hash.get(sum));
} else {
hash.put(sum, i);
}
}
return ret;
}
}
复杂度分析:
- 时间复杂度:O(n)
- 空间复杂度:O(n)
032 矩阵区域和
题目描述:
给定一个 m x n 的矩阵 mat 和一个整数 k,返回一个新的矩阵 answer,其中 answer[i][j] 等于原矩阵中所有满足 |r - i| <= k 且 |c - j| <= k 的元素 mat[r][c] 的和。
思路分析
这是一道典型的二维前缀和应用题。直接对每个点暴力累加周围区域会导致 O(mnk²) 的时间复杂度。利用二维前缀和可以将单次查询优化至 O(1)。
定义 dp[i][j] 为以 (0, 0) 为左上角,(i-1, j-1) 为右下角的矩形区域内所有元素的和。注意这里通常将 dp 数组开大一行一列,方便处理边界情况,避免判断下标是否越界。
递推公式如下:
dp[i][j] = dp[i-1][j] + dp[i][j-1] - dp[i-1][j-1] + mat[i-1][j-1]
对于结果矩阵中的每个点 (i, j),我们需要计算其周围 k 范围内的和。这对应于原矩阵中的一个矩形区域,我们需要确定该区域的左上角 (x1, y1) 和右下角 (x2, y2) 坐标。由于范围可能超出矩阵边界,需要进行裁剪:
- 左上角:
x1 = max(0, i - k),y1 = max(0, j - k) - 右下角:
x2 = min(m - 1, i + k),y2 = min(n - 1, j + k)
代入二维前缀和公式即可得到区域和。注意 dp 数组的下标映射关系,因为 dp 多了一行一列,实际坐标需要 +1。
代码实现
C++ 实现
class Solution {
public:
vector<vector<int>> matrixBlockSum(vector<vector<int>>& mat, int k) {
int m = mat.size(), n = mat[0].size();
// 预处理前缀和矩阵,大小为 (m+1) x (n+1)
vector<vector<int>> dp(m + 1, vector<int>(n + 1, 0));
// 填充 dp 表
for (int i = 1; i <= m; ++i) {
for (int j = 1; j <= n; ++j) {
dp[i][j] = dp[i - 1][j] + dp[i][j - 1] - dp[i - 1][j - 1] + mat[i - 1][j - 1];
}
}
vector<vector<int>> ret(m, vector<int>(n));
// 计算每个点的区域和
for (int i = 0; i < m; ++i) {
for (int j = 0; j < n; ++j) {
// 计算矩形区域的四个角在 dp 表中的坐标
int x1 = max(0, i - k) + 1, y1 = max(0, j - k) + 1;
int x2 = min(m - 1, i + k) + 1, y2 = min(n - 1, j + k) + 1;
// 利用容斥原理计算区域和
ret[i][j] = dp[x2][y2] - dp[x1 - 1][y2] - dp[x2][y1 - 1] + dp[x1 - 1][y1 - 1];
}
}
return ret;
}
};
复杂度分析:
- 时间复杂度:O(m * n),预处理和查询各需遍历一次矩阵。
- 空间复杂度:O(m * n),用于存储前缀和矩阵。
Java 实现
class Solution {
public int[][] matrixBlockSum(int[][] mat, int k) {
int m = mat.length, n = mat[0].length;
// 1、预处理前缀和矩阵
int[][] dp = new int[m + 1][n + 1];
for (int i = 1; i <= m; i++) {
for (int j = 1; j <= n; j++) {
dp[i][j] = dp[i - 1][j] + dp[i][j - 1] - dp[i - 1][j - 1] + mat[i - 1][j - 1];
}
}
// 2、使用前缀和矩阵计算结果
int[][] ret = new int[m][n];
for (int i = 0; i < m; i++) {
for (int j = 0; j < n; j++) {
int x1 = Math.max(0, i - k) + 1, y1 = Math.max(0, j - k) + 1;
int x2 = Math.min(m - 1, i + k) + 1, y2 = Math.min(n - 1, j + k) + 1;
ret[i][j] = dp[x2][y2] - dp[x1 - 1][y2] - dp[x2][y1 - 1] + dp[x1 - 1][y1 - 1];
}
}
return ret;
}
}
复杂度分析:
- 时间复杂度:O(m * n)
- 空间复杂度:O(m * n)
总结
这两道题展示了前缀和在不同维度下的应用。一维场景下结合哈希表可以解决特定条件下的最长子数组问题;二维场景下则通过预处理矩阵实现快速区域查询。建议读者在理解公式推导的基础上,动手画图模拟边界情况,特别是坐标转换和数组下标偏移的细节,这对面试中的边界测试至关重要。

