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.

python
# Pseudo-code for Brute Force
count = 0
for r in range(rows):
    for c in range(cols):
        if matrix[r][c] == 1:
            max_side = min(rows - r, cols - c)
            for k in range(1, max_side + 1):
                if is_all_ones(r, c, k): # This check takes O(k^2)
                    count += 1
                else:
                    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.

Core Insight for Count Square Submatrices with All Ones

The key intuition is to change the definition of what we are counting. Instead of asking "Does a square of size $k$ start here?", we ask "What is the maximum size of a square that ends at this cell (bottom-right corner)?"

Let dp[i][j] represent the side length of the largest square whose bottom-right corner is at (i, j).

If matrix[i][j] is 1, the size of the square ending at (i, j) is limited by its three neighbors:

1. The square ending directly above it (i-1, j).
2. The square ending directly to the left (i, j-1).
3. The square ending diagonally top-left (i-1, j-1).

The Constraint: To form a square of size $k$ at (i, j), the neighbors (Top, Left, and Top-Left) must all support squares of size at least $k-1$. If any neighbor has a smaller square, the square at (i, j) is bottlenecked by that minimum size.

Therefore, the recurrence relation is: $dp[i][j] = \min(dp[i-1][j], dp[i][j-1], dp[i-1][j-1]) + 1$

Visual Description: Imagine building a square block by block. To place a $2 \times 2$ block ending at $(i, j)$, you need a $1 \times 1$ block at the top, left, and top-left positions. To place a $3 \times 3$ block, those neighbors must themselves be the corners of $2 \times 2$ blocks. The smallest neighbor dictates the maximum extension possible.

Crucially, the value dp[i][j] also tells us exactly how many squares end at (i, j). If dp[i][j] = 3, it means a $3 \times 3$, a $2 \times 2$, and a $1 \times 1$ square all end there. Thus, the sum of all values in the DP table is the answer.

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$).

Reference Implementation

C++ Solution for LeetCode 1277

cpp
class Solution {
public:
    int countSquares(vector<vector<int>>& matrix) {
        int m = matrix.size();
        int n = matrix[0].size();
        int totalSquares = 0;
        
        for (int i = 0; i < m; ++i) {
            for (int j = 0; j < n; ++j) {
                // Only process if the current cell is a potential corner of a square
                if (matrix[i][j] == 1) {
                    // If we are not on the boundary, apply DP transition
                    if (i > 0 && j > 0) {
                        int top = matrix[i-1][j];
                        int left = matrix[i][j-1];
                        int diag = matrix[i-1][j-1];
                        
                        // Update current cell with size of largest square ending here
                        matrix[i][j] = min({top, left, diag}) + 1;
                    }
                    // Accumulate the size (which equals the count of squares ending here)
                    totalSquares += matrix[i][j];
                }
            }
        }
        
        return totalSquares;
    }
};

Java Solution for LeetCode 1277

java
class Solution {
    public int countSquares(int[][] matrix) {
        int m = matrix.length;
        int n = matrix[0].length;
        int totalSquares = 0;
        
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                // If the cell is 1, it contributes to the count
                if (matrix[i][j] == 1) {
                    // Check boundaries to avoid IndexOutOfBounds
                    if (i > 0 && j > 0) {
                        int top = matrix[i-1][j];
                        int left = matrix[i][j-1];
                        int diag = matrix[i-1][j-1];
                        
                        // Update the matrix in-place with the recurrence relation
                        matrix[i][j] = Math.min(Math.min(top, left), diag) + 1;
                    }
                    // Add the value of the current cell to total
                    totalSquares += matrix[i][j];
                }
            }
        }
        
        return totalSquares;
    }
}

Python Solution for LeetCode 1277

python
class Solution:
    def countSquares(self, matrix: List[List[int]]) -> int:
        rows = len(matrix)
        cols = len(matrix[0])
        total_squares = 0
        
        for i in range(rows):
            for j in range(cols):
                # We only care if the cell initiates or continues a square
                if matrix[i][j] == 1:
                    # If inside the grid boundaries (not first row/col)
                    if i > 0 and j > 0:
                        # Find the bottleneck among the three neighbors
                        matrix[i][j] = min(matrix[i-1][j], 
                                           matrix[i][j-1], 
                                           matrix[i-1][j-1]) + 1
                    
                    # Add the size of the square ending at (i, j) to the total
                    total_squares += matrix[i][j]
                    
        return total_squares

Why is my answer just the count of 1s?

• Mistake: Forgetting to update the DP table (or matrix) and just summing the original 1s.
• Why: Failing to accumulate state. The algorithm needs to "remember" that a larger square was formed so subsequent cells can build upon it.
• Consequence: You only count squares of size $1 \times 1$.