Problem Statement in English

You’re given a 2D grid of integers and an integer k. You have to count the number of submatrices that have their top-left element less than or equal to k and the sum of all elements in the submatrix is less than or equal to k.

Approach

To solve this problem, we can use a 2D prefix sum matrix. The prefix sum matrix allows us to quickly calculate the sum of any submatrix in constant time after an initial preprocessing step.

To do this, read up my post 304. Range Sum Query 2D - Immutable where I have explained how to build a 2D prefix sum matrix and query it.

And we’re done!

Solution in Python


class Solution:
    def countSubmatrices(self, grid: List[List[int]], k: int) -> int:
        M = len(grid)
        N = len(grid[0])

        matrix = [[0 for _ in range(N)] for _ in range(M)]
        matrix[0][0] = grid[0][0]

        # build the first row
        for j in range(1, N):
            matrix[0][j] = matrix[0][j - 1] + grid[0][j] 

        # build the first col
        for i in range(1, M):
            matrix[i][0] = matrix[i - 1][0] + grid[i][0]

        # build 2D prefix sum
        for i in range(1, M):
            for j in range(1, N):
                # grid this cell + matrix cell above + matrix cell left - matrix cell diag left
                matrix[i][j] = grid[i][j] + matrix[i - 1][j] + matrix[i][j - 1] - matrix[i - 1][j - 1]
        
        res = 0
        # query the 2D prefix sum
        for i in range(M):
            for j in range(N):
                if matrix[i][j] <= k:
                    res += 1

        return res

Complexity

  • Time: $O(M \times N)$
    Where $M$ and $N$ are the number of rows and columns in the grid respectively. We build the prefix sum matrix in $O(M \times N)$ time and then we query it in $O(M \times N)$ time.

  • Space: $O(M \times N)$
    We use an additional matrix to store the prefix sums. But we can optimize this to $O(1)$ space by modifying the input grid to store the prefix sums instead of using an additional matrix.

And we are done.