LeetCode 3212 - Count Submatrices With Equal Frequency of X and Y

The problem asks us to count submatrices within a given 2D character matrix grid that satisfy three conditions. A submatrix is defined by a contiguous rectangle within the grid, and the submatrix must include the top-left cell grid[0][0].

leetcodemediumarraymatrixprefix-sum

LeetCode Problem 3212

Difficulty: 🟡 Medium
Topics: Array, Matrix, Prefix Sum

Solution

Problem Understanding

The problem asks us to count submatrices within a given 2D character matrix grid that satisfy three conditions. A submatrix is defined by a contiguous rectangle within the grid, and the submatrix must include the top-left cell grid[0][0]. The first condition is that the submatrix contains equal numbers of 'X' and 'Y'. The second condition is that there is at least one 'X' in the submatrix. The input grid can also contain the character '.', which represents a neutral cell that does not affect counts of 'X' or 'Y'.

The output is a single integer representing the number of valid submatrices that meet these criteria. Constraints indicate that the grid can be as large as 1000x1000, which rules out naive O(n^6) solutions that enumerate every possible submatrix and count characters individually.

Important edge cases include grids without any 'X' (output must be 0), grids without any 'Y' (output must be 0 unless equal counts of 'X' and 'Y' can exist), and grids where all cells are '.'. Another subtle point is that every valid submatrix must include grid[0][0], which restricts the set of submatrices significantly.

Approaches

Brute Force

The brute-force approach would enumerate all possible submatrices containing grid[0][0]. For each candidate submatrix, count the occurrences of 'X' and 'Y' and check whether they are equal and that there is at least one 'X'.

This is correct, as it checks every possible submatrix, but it is too slow because for an n x m grid, there are O(n^2 * m^2) submatrices. Counting 'X' and 'Y' in each submatrix requires another O(n * m) operation, yielding O(n^3 * m^3) overall, which is infeasible for n, m up to 1000.

Optimal Approach

The key insight is that we can transform the problem into a prefix sum problem using a value mapping for 'X' as +1, 'Y' as -1, and '.' as 0. The goal of having equal counts of 'X' and 'Y' is equivalent to finding submatrices whose summed values equal 0. The requirement of at least one 'X' can be enforced separately.

We can then reduce the 2D submatrix sum problem to a series of 1D subarray sum problems using prefix sums along rows and iterating over all possible column ranges. For each column range, we compute the running sum of the transformed values across rows and use a hashmap to count the number of contiguous row segments whose sum is 0, similar to the classic "subarray sum equals K" problem.

This drastically reduces complexity because it leverages cumulative sums and hashmaps to count valid submatrices in O(n * m^2) time rather than O(n^3 * m^3).

Approach Time Complexity Space Complexity Notes
Brute Force O(n^3 * m^3) O(1) Enumerate all submatrices and count 'X' and 'Y'
Optimal O(n * m^2) O(n) Use prefix sums and hashmap to count zero-sum row ranges

Algorithm Walkthrough

  1. Transform the input grid to a numeric grid where 'X' = 1, 'Y' = -1, and '.' = 0. This allows us to convert the equal count condition into a zero-sum problem.

  2. Compute prefix sums along rows to allow O(1) range sum queries for any column range.

  3. Iterate over all pairs of columns (left, right) representing the vertical boundaries of submatrices. For each pair:

  4. For each row, compute the cumulative sum of values in the column range [left, right].

  5. Track the running cumulative sum of these row sums and store counts in a hashmap.

  6. Increment the answer whenever a zero cumulative sum is encountered in a contiguous segment, ensuring that the submatrix sum is zero.

  7. For each candidate submatrix, verify that it contains at least one 'X' within the column range. This can be done by separately tracking 'X' counts for each row segment.

  8. Return the total count of valid submatrices.

Why it works: Mapping 'X' to 1 and 'Y' to -1 reduces the equal count condition to a sum of zero, which can be efficiently detected using prefix sums and a hashmap. This ensures correctness while maintaining performance.

Python Solution

from typing import List
from collections import defaultdict

class Solution:
    def numberOfSubmatrices(self, grid: List[List[str]]) -> int:
        if not grid or not grid[0]:
            return 0
        n, m = len(grid), len(grid[0])
        # Convert grid to numeric values: X=1, Y=-1, .=0
        num_grid = [[1 if cell=='X' else -1 if cell=='Y' else 0 for cell in row] for row in grid]
        x_grid = [[1 if cell=='X' else 0 for cell in row] for row in grid]

        # Prefix sum for each row
        for r in range(n):
            for c in range(1, m):
                num_grid[r][c] += num_grid[r][c-1]
                x_grid[r][c] += x_grid[r][c-1]

        ans = 0
        for left in range(m):
            for right in range(left, m):
                cum_sum = 0
                x_count = 0
                counter = defaultdict(int)
                counter[0] = 1
                for r in range(n):
                    val = num_grid[r][right] - (num_grid[r][left-1] if left > 0 else 0)
                    x_val = x_grid[r][right] - (x_grid[r][left-1] if left > 0 else 0)
                    cum_sum += val
                    x_count += x_val
                    if x_count > 0:
                        ans += counter[cum_sum]
                    counter[cum_sum] += 1
        return ans

This solution first maps the grid to numeric arrays for quick computation. Row-wise prefix sums allow efficient column range queries. For each column pair, a hashmap counts cumulative sums of row segments, and we ensure that submatrices contain at least one 'X'.

Go Solution

func numberOfSubmatrices(grid [][]byte) int {
    n, m := len(grid), len(grid[0])
    numGrid := make([][]int, n)
    xGrid := make([][]int, n)
    for i := 0; i < n; i++ {
        numGrid[i] = make([]int, m)
        xGrid[i] = make([]int, m)
        for j := 0; j < m; j++ {
            if grid[i][j] == 'X' {
                numGrid[i][j] = 1
                xGrid[i][j] = 1
            } else if grid[i][j] == 'Y' {
                numGrid[i][j] = -1
            }
        }
        for j := 1; j < m; j++ {
            numGrid[i][j] += numGrid[i][j-1]
            xGrid[i][j] += xGrid[i][j-1]
        }
    }

    ans := 0
    for left := 0; left < m; left++ {
        for right := left; right < m; right++ {
            cumSum := 0
            xCount := 0
            counter := map[int]int{0: 1}
            for r := 0; r < n; r++ {
                val := numGrid[r][right]
                xVal := xGrid[r][right]
                if left > 0 {
                    val -= numGrid[r][left-1]
                    xVal -= xGrid[r][left-1]
                }
                cumSum += val
                xCount += xVal
                if xCount > 0 {
                    ans += counter[cumSum]
                }
                counter[cumSum]++
            }
        }
    }
    return ans
}

The Go implementation follows the same logic as Python, using slices instead of lists and maps for hash counting. It handles column prefix sums carefully and ensures proper initialization of the counter map.

Worked Examples

Example 1: grid = [["X","Y","."],["Y",".","."]]

Step through column ranges:

  1. left=0, right=0: row sums [1, -1], cumulative sums [1,0], count of X in segments [1,1] → valid segments detected.
  2. left=0, right=1: row sums [0, -1], cumulative sums [0,-1], count of X [1,1] → valid segments detected.
  3. left=0, right=2: row sums [0, -1], cumulative sums [0,-1] → valid segments.
  4. Other column ranges yield no new valid submatrices.

Total count = 3.

Example 2: grid = `[["X","X"],["X","