Count Square Submatrices with All Ones

LeetCode 1277: Count Square Submatrices with All Ones Solution & Explanation

Problem Overview

TL;DR: Use 2D Dynamic Programming to compute the size of the largest square ending at every cell, then sum these sizes to get the total count.

The LeetCode 1277 problem, "Count Square Submatrices with All Ones," asks us to analyze a binary matrix (containing only 0s and 1s) and determine the total number of square submatrices that consist entirely of 1s. A square submatrix can be of any size ($1 \times 1$, $2 \times 2$, etc.), provided it fits within the grid boundaries.

This is a popular interview question because it requires optimizing a geometric search problem into a linear-time pass using state accumulation.

Brute Force Approach for Count Square Submatrices with All Ones

A naive approach involves iterating through every cell in the matrix and treating it as the top-left corner of a potential square. For each corner, we attempt to expand the square size ($k=1, k=2, \dots$) as long as all cells within that boundary are 1s.

Naive Algorithm:

1. Iterate through every cell (i, j) in the matrix.
2. For each cell, iterate through all possible side lengths k.
3. Check if the submatrix from (i, j) to (i+k-1, j+k-1) contains only ones.
4. If yes, increment the count. If no, stop expanding for this corner.

1# Pseudo-code for Brute Force
2count = 0
3for r in range(rows):
4    for c in range(cols):
5        if matrix[r][c] == 1:
6            max_side = min(rows - r, cols - c)
7            for k in range(1, max_side + 1):
8                if is_all_ones(r, c, k): # This check takes O(k^2)
9                    count += 1
10                else:
11                    break

Why it fails: The time complexity of this approach is roughly $O(m \cdot n \cdot \min(m,n)^2)$. For a $300 \times 300$ matrix, this results in roughly $8 \times 10^9$ operations in the worst case, leading to a Time Limit Exceeded (TLE) error. The fundamental inefficiency is the redundant checking of cells; knowing that a $3 \times 3$ square exists implies a $2 \times 2$ square also exists, but the brute force re-verifies this from scratch.

Algorithm Strategy: Dynamic Programming (DP) Patterns

We will utilize a tabulation strategy (bottom-up DP) to solve this efficiently.

1. State Initialization: We can create a separate DP table of the same size as the matrix, or modify the input matrix in-place to save space (if allowed).
2. Boundary Handling:
   • The first row (i=0) and first column (j=0) cannot look back at i-1 or j-1.
   • For these cells, dp[i][j] is simply equal to matrix[i][j]. A cell in the first row can only form a square of size 1.
3. Iterative Transition:
   • Iterate through the matrix starting from (1, 1).
   • If matrix[i][j] == 1, apply the recurrence relation: min(top, left, diag) + 1.
   • If matrix[i][j] == 0, dp[i][j] is 0 (no square ends there).
4. Accumulation: Keep a running total of all values calculated in the DP table.

Execution Flow

1. Initialize a variable total_squares = 0.
2. Iterate through the grid with row index i from 0 to m-1 and column index j from 0 to n-1.
3. Check Condition: If matrix[i][j] == 1:
   • Boundary Case: If i == 0 or j == 0, the max square size is 1. We just add 1 to total_squares. (In an in-place modification, we leave the 1 as is).
   • Internal Case: If i > 0 and j > 0, calculate the value: val = min(matrix[i-1][j], matrix[i][j-1], matrix[i-1][j-1]) + 1.
   • Update matrix[i][j] (or the DP table) with val.
   • Add val to total_squares.
4. If matrix[i][j] == 0, do nothing (it contributes 0).
5. Return total_squares.

Proof of Correctness

The correctness relies on the geometric property of squares. A square of size $S \times S$ ending at $(i, j)$ covers the region from $(i-S+1, j-S+1)$ to $(i, j)$. For this to exist, the submatrix defined by $(i-1, j)$ must support a square of size $S-1$, as must $(i, j-1)$ and $(i-1, j-1)$.

By taking the minimum of these three neighbors and adding 1, we guarantee that we extend the square only as far as the "weakest" link allows, ensuring no 0s are included. Summing these values correctly counts all squares because a max-square of size $k$ implies exactly $k$ valid squares ending at that specific coordinate (sizes $1, 2, ..., k$).

Pattern Reuse Notes

The logic used in LeetCode 1277 Solution is a classic application of DP on grids. This specific "look back at neighbors" pattern appears in several other problems:

• LeetCode 62: Unique Paths - Sums the paths from Top and Left.
• LeetCode 63: Unique Paths II - Same as above but handles obstacles (0s).
• LeetCode 64: Minimum Path Sum - Takes the min of Top and Left plus current cost.
• LeetCode 221: Maximal Square - Identical logic to LeetCode 1277, but asks for the maximum side length rather than the total count.

Reference Implementation

C++ Solution for LeetCode 1277

1class Solution {
2public:
3    int countSquares(vector<vector<int>>& matrix) {
4        int m = matrix.size();
5        int n = matrix[0].size();
6        int totalSquares = 0;
7        
8        for (int i = 0; i < m; ++i) {
9            for (int j = 0; j < n; ++j) {
10                // Only process if the current cell is a potential corner of a square
11                if (matrix[i][j] == 1) {
12                    // If we are not on the boundary, apply DP transition
13                    if (i > 0 && j > 0) {
14                        int top = matrix[i-1][j];
15                        int left = matrix[i][j-1];
16                        int diag = matrix[i-1][j-1];
17                        
18                        // Update current cell with size of largest square ending here
19                        matrix[i][j] = min({top, left, diag}) + 1;
20                    }
21                    // Accumulate the size (which equals the count of squares ending here)
22                    totalSquares += matrix[i][j];
23                }
24            }
25        }
26        
27        return totalSquares;
28    }
29};

Java Solution for LeetCode 1277

1class Solution {
2    public int countSquares(int[][] matrix) {
3        int m = matrix.length;
4        int n = matrix[0].length;
5        int totalSquares = 0;
6        
7        for (int i = 0; i < m; i++) {
8            for (int j = 0; j < n; j++) {
9                // If the cell is 1, it contributes to the count
10                if (matrix[i][j] == 1) {
11                    // Check boundaries to avoid IndexOutOfBounds
12                    if (i > 0 && j > 0) {
13                        int top = matrix[i-1][j];
14                        int left = matrix[i][j-1];
15                        int diag = matrix[i-1][j-1];
16                        
17                        // Update the matrix in-place with the recurrence relation
18                        matrix[i][j] = Math.min(Math.min(top, left), diag) + 1;
19                    }
20                    // Add the value of the current cell to total
21                    totalSquares += matrix[i][j];
22                }
23            }
24        }
25        
26        return totalSquares;
27    }
28}

Python Solution for LeetCode 1277

1class Solution:
2    def countSquares(self, matrix: List[List[int]]) -> int:
3        rows = len(matrix)
4        cols = len(matrix[0])
5        total_squares = 0
6        
7        for i in range(rows):
8            for j in range(cols):
9                # We only care if the cell initiates or continues a square
10                if matrix[i][j] == 1:
11                    # If inside the grid boundaries (not first row/col)
12                    if i > 0 and j > 0:
13                        # Find the bottleneck among the three neighbors
14                        matrix[i][j] = min(matrix[i-1][j], 
15                                           matrix[i][j-1], 
16                                           matrix[i-1][j-1]) + 1
17                    
18                    # Add the size of the square ending at (i, j) to the total
19                    total_squares += matrix[i][j]
20                    
21        return total_squares