Editorial

Solution walkthrough

Review the explanation, then compare the reference implementation.

Video explanation

Solution

Intuition

An axis-aligned submatrix containing the grid's top-left cell must start at that cell. Its bottom-right corner uniquely identifies it, so there is just one candidate per grid cell. A two-dimensional prefix sum gives each candidate's total.

Approach

Build a prefix-sum matrix where ps[i][j] is the sum from (0, 0) through (i, j). Initialize the top-left value, then accumulate the first row and first column. For every interior cell, add the prefix above and the prefix to the left, subtract their doubly counted overlap, and add the current grid value.

Whenever a prefix sum is at most k, increment the answer. The comparison is inclusive, as required by the problem statement. Since each candidate corresponds to exactly one bottom-right corner, counting these qualifying prefixes counts every valid submatrix once.

class Solution:
    def countSubmatrices(self, grid: List[List[int]], k: int) -> int:
        n = len(grid)
        m = len(grid[0])
        ps = [[0] * m for _ in range(n)]
        ps[0][0] = grid[0][0]
        res = 1 if ps[0][0] <= k else 0

        for j in range(1, m):
            ps[0][j] = grid[0][j] + ps[0][j - 1]
            if ps[0][j] <= k:
                res += 1

        for i in range(1, n):
            ps[i][0] = grid[i][0] + ps[i - 1][0]
            if ps[i][0] <= k:
                res += 1

        for i in range(1, n):
            for j in range(1, m):
                ps[i][j] = grid[i][j] + (
                    ps[i - 1][j] + ps[i][j - 1] - ps[i - 1][j - 1]
                )
                if ps[i][j] <= k:
                    res += 1

        return res

Complexity analysis

Let r and c be the grid dimensions. Time: O(r * c). Auxiliary space: O(r * c) for the two-dimensional prefix-sum matrix.