【数据结构-二维前缀和】【含空间优化】力扣3070. 元素和小于等于 k 的子矩阵的数目

给你一个下标从 0 开始的整数矩阵 grid 和一个整数 k。

返回包含 grid 左上角元素、元素和小于或等于 k 的
子矩阵
的数目。

示例 1：

输入：grid = [[7,6,3],[6,6,1]], k = 18
输出：4
解释：如上图所示，只有 4 个子矩阵满足：包含 grid 的左上角元素，并且元素和小于或等于 18 。

示例 2：

输入：grid = [[7,2,9],[1,5,0],[2,6,6]], k = 20
输出：6
解释：如上图所示，只有 6 个子矩阵满足：包含 grid 的左上角元素，并且元素和小于或等于 20 。

法一：二维前缀和

class Solution {
public:
    int countSubmatrices(vector<vector<int>>& grid, int k) {
        int m = grid.size(), n = grid[0].size();
        vector<vector<int>> sums(m+1, vector<int>(n+1));
        int count = 0;
        for(int i = 0;i < m; i++){
            for(int j = 0; j < n; j++){
                sums[i+1][j+1] = sums[i+1][j] + sums[i][j+1] - sums[i][j] + grid[i][j];
                if(sums[i+1][j+1]>k){
                    break;
                }else{
                    count++;
                }
            }
        }
        return count;
    }
};

时间复杂度：O(mn)，其中 m 和 n 分别为 grid 的行数和列数。
空间复杂度：O(mn)。

这种方法就是传统的前缀和做法，在这道题中，如果group[i][j]小于k，则继续遍历，然后count增加，如果group[i][j]大于k，说明不需要继续遍历这一行接下来列的前缀和了，直接开始计算新一行每一列的前缀和。

空间优化：维护列

class Solution {
public:
    int countSubmatrices(vector<vector<int>>& grid, int k) {
        int m = grid.size(), n = grid[0].size();
        int count = 0;
        vector<int> col_sum(n);
        for(auto &row : grid){
            int s = 0;
            for(int j = 0;j < n; j++){
                col_sum[j] += row[j];
                s += col_sum[j];
                if(s > k) break;
                count++;
            }
        }
        return count;
    }
};

时间复杂度：O(mn)，其中 m 和 n 分别为 grid 的行数和列数。
空间复杂度：O(n)。

定义一个一维数组col_sum来记录每一列的元素和，为什么可以这样做呢？因为矩阵中的元素都大于0，当计算某一行的某一列前缀和的时候，如果发现前缀和超过k，那么在接下来的行的计算中，就不会再计算超过这一列的前缀和。

这时候我们定义一个整型s来记录每一行某个位置的前缀和，并且s在重新计算新一行的时候会重置为0。col_sum[j]在遍历行的时候，也会依次计算该列的元素和，要计算前缀和s，实际上就是目前该列的元素和逐列累加，当发现s大于k的时候，那么就停止该行接下来的前缀和，继续计算下一行第一列开始的前缀和，继续遍历行的某个位置的前缀和。

最后遍历完所有行后，返回count即可。