LeetCode 3070. 元素和小于等于 k 的子矩阵数目
题目描述
给你一个大小为 m x n 的整数矩阵 grid 和一个整数 k。你需要找出 grid 中所有以左上角 (0,0) 为起始点的子矩阵,并统计这些子矩阵中元素和不超过 k 的个数。
注意 :子矩阵必须包含 (0,0) 这个格子,即子矩阵的左上角固定为原点,右下角可以是任意坐标 (i,j)(0 <= i < m, 0 <= j < n)。
思路分析
本题的暴力做法是枚举所有可能的子矩阵,并计算其和,但时间复杂度会达到 O(m²n²),不可取。由于子矩阵的左上角固定,我们可以利用二维前缀和 来快速得到任意从 (0,0) 到 (i,j) 的子矩阵的和。
二维前缀和
定义 s[i+1][j+1] 表示原矩阵中从 (0,0) 到 (i,j) 的子矩阵的元素和。这样 s[i][j] 就对应了以 (i-1,j-1) 为右下角的子矩阵的和(左上角为 (0,0))。
前缀和的递推公式(容斥原理):
s[i+1][j+1] = s[i+1][j] + s[i][j+1] - s[i][j] + grid[i][j]
其中 s[0][*] 和 s[*][0] 均为 0,方便边界处理。
统计答案
构建完前缀和数组 s 后,我们只需要遍历所有可能的右下角坐标 (i,j)(i 从 1 到 m,j 从 1 到 n),检查 s[i][j] <= k 是否成立,如果成立则答案加一。
代码实现
cpp
class Solution {
public:
int countSubmatrices(vector<vector<int>>& grid, int k) {
int m = grid.size(), n = grid[0].size();
// 二维前缀和数组,大小为 (m+1) x (n+1)
vector<vector<int>> s(m + 1, vector<int>(n + 1, 0));
// 构建前缀和
for (int i = 0; i < m; ++i) {
for (int j = 0; j < n; ++j) {
s[i + 1][j + 1] = s[i + 1][j] + s[i][j + 1] - s[i][j] + grid[i][j];
}
}
int ans = 0;
// 枚举右下角
for (int i = 1; i <= m; ++i) {
for (int j = 1; j <= n; ++j) {
if (s[i][j] <= k) ++ans;
}
}
return ans;
}
};
复杂度分析
- 时间复杂度:O(m × n)。构建前缀和需要遍历所有格子,统计答案也需要遍历所有可能的右下角,因此总时间复杂度为 O(mn)。
- 空间复杂度:O(m × n),用于存储二维前缀和数组。
总结
本题的核心是二维前缀和的应用。由于子矩阵的左上角固定为原点,前缀和数组可以直接对应每个可能子矩阵的和,从而将统计过程简化为一次遍历。掌握二维前缀和可以高效解决类似矩阵区域和的问题。