Problem Statement in English

You’re given a 2D grid of characters. You have to count the number of submatrices that have an equal frequency of ‘X’ and ‘Y’. The submatrix must include the top left cell of the grid, and must have atleast one ‘X’.

Return the count of such submatrices.

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 numberOfSubmatrices(self, grid: List[List[str]]) -> int:
        M = len(grid)
        N = len(grid[0])

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

        def repr(i, j):
            c = grid[i][j]
            if c == "X":
                return [1, 0]
            elif c == "Y":
                return [0, 1]
            return [0, 0]

        def add(*vectors):
            result = [0, 0]

            for v in vectors:
                result[0] += v[0]
                result[1] += v[1]

            return result

        def subtract(x, y):
            return [x[0] - y[0], x[1] - y[1]]

        matrix[0][0] = repr(0, 0)

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

        # build the first col
        for i in range(1, M):
            matrix[i][0] = add(matrix[i - 1][0], repr(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] = subtract(
                    add(
                        repr(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][0] and matrix[i][j][0] == matrix[i][j][1]:
                    res += 1

        return res

Complexity

  • Time: $O(M \times N)$
    Where $M$ and $N$ are the dimensions of the grid. We build a 2D prefix sum in $O(M \times N)$ time and then query it in $O(M \times N)$ time.

  • Space: $O(M \times N)$
    We build a 2D prefix sum which takes $O(M \times N)$ space.

And we are done.