LeetCode 3070. 元素和小于等于 k 的子矩阵的数目

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 。

提示：

• m == grid.length
• n == grid[i].length
• 1 <= n, m <= 1000
• 0 <= grid[i][j] <= 1000
• 1 <= k <= 10^9

解法1：二维前缀和 + 双重循环

Java版：

class Solution {
    public int countSubmatrices(int[][] grid, int k) {
        int m = grid.length;
        int n = grid[0].length;
        int[][] presum = new int[m + 1][n + 1];
        int ans = 0;
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                presum[i + 1][j + 1] = presum[i + 1][j] + presum[i][j + 1] - presum[i][j] + grid[i][j];
                if (presum[i + 1][j + 1] <= k) {
                    ans++;
                }
            }
        }
        return ans;
    } 
}

Python3版：

class Solution:
    def countSubmatrices(self, grid: List[List[int]], k: int) -> int:
        m = len(grid)
        n = len(grid[0])
        presum = [[0] * (n + 1) for _ in range(m + 1)]
        ans = 0
        for i in range(m):
            for j in range(n):
                presum[i + 1][j + 1] = presum[i + 1][j] + presum[i][j + 1] - presum[i][j] + grid[i][j]
                if presum[i + 1][j + 1] <= k:
                    ans += 1
        return ans

复杂度分析

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

解法2：维护每一列的元素和

遍历每一行，同时用一个长为 n 的数组 colsum 维护每一列的元素和。

遍历当前行时，一边更新 colsum[j]，一边累加 colsum[j] 到变量 s 中。

如果 s≤k 则把答案加一，否则可以退出循环（因为矩阵元素都非负）。

Java版：

class Solution {
    public int countSubmatrices(int[][] grid, int k) {
        int m = grid.length;
        int n = grid[0].length;
        int[] colsum = new int[n];
        int ans = 0;
        for (int i = 0; i < m; i++) {
            int s = 0;
            for (int j = 0; j < n; j++) {
                colsum[j] += grid[i][j];
                s += colsum[j];
                if (s <= k) {
                    ans++;
                } else {
                    break;
                }
            }
        } 
        return ans;
    }
}

Python3版：

class Solution:
    def countSubmatrices(self, grid: List[List[int]], k: int) -> int:
        m = len(grid)
        n = len(grid[0])
        colsum = [0] * n
        ans = 0
        for i in range(m):
            s = 0
            for j in range(n):
                colsum[j] += grid[i][j]
                s += colsum[j]
                if s <= k:
                    ans += 1
                else:
                    break
        return ans

复杂度分析

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

解法3：原地修改数组第一行

把每列元素和保存到 grid 的第一行，可以做到 O(1) 额外空间。

Java版：

class Solution {
    public int countSubmatrices(int[][] grid, int k) {
        int m = grid.length;
        int n = grid[0].length;
        int ans = 0;
        for (int i = 0; i < m; i++) {
            int s = 0;
            for (int j = 0; j < n; j++) {
                if (i > 0) {
                    grid[0][j] += grid[i][j];
                }
                s += grid[0][j];
                if (s <= k) {
                    ans++;
                } else {
                    break;
                }
            }
        } 
        return ans;
    }
}

Python3版：

class Solution:
    def countSubmatrices(self, grid: List[List[int]], k: int) -> int:
        m = len(grid)
        n = len(grid[0])
        ans = 0
        for i in range(m):
            s = 0
            for j in range(n):
                if i > 0:
                    grid[0][j] += grid[i][j]
                s += grid[0][j]
                if s <= k:
                    ans += 1
                else:
                    break
        return ans

复杂度分析

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