DP 56. Count Square Submatrices with All Ones | DP on Rectangles

1. Problem Statement & Overview

Given an $n \times m$ binary matrix filled with 0s and 1s, the objective is to count and return the total number of square submatrices that contain only 1s. A square requires equal length and breadth. The problem can be solved across multiple sizes (e.g., size 1, size 2, size 3, etc.) simultaneously by evaluating how larger squares can be formed using smaller neighboring substructures.

2. Why Do We Need This? (The Limitations of Brute Force)

A naive brute-force approach requires standing at every single cell (i, j) and attempting to expand outward in all possible directions to check if a valid square submatrix can be formed.

• The Problem: Checking every potential square by validating all interior elements leads to immense computational redundancy.
• Complexity: Expanding from every cell and verifying matrix elements takes $O(n^2 \cdot m^2)$ or worse, making the code extremely complex and prone to Time Limit Exceeded (TLE) errors in interviews and online judges.
• The Solution: Dynamic Programming (DP) optimizes this evaluation by storing intermediate square counts, reducing the time complexity to linear space and time relative to the matrix dimensions.

3. Core Mental Model & DP Definition

To solve this efficiently, we construct a 2D dynamic programming table (dp[][]) of the exact same dimensions as the given matrix (n x m).

• Definition of dp[i][j]: dp[i][j] represents the maximum side length of the square submatrix whose right-bottom corner is located at cell (i, j).
• Key Insight: If dp[i][j] = k, it implies there are $k$ squares of sizes $1, 2, \dots, k$ whose bottom-right corner terminates precisely at coordinates (i, j). Therefore, summing up all values across the entire dp[][] table at the end yields the total count of square submatrices in the matrix.

4. Naive vs. Better Approach

5. Step-by-Step Evolution & Recurrence Relation

If the current cell in the original matrix contains a 0, it cannot be part of any square of ones, so dp[i][j] = 0. However, if the current cell contains a 1, its capacity to form a larger square depends strictly on the minimum square sizes already established by its immediate neighbors:

1. The Top Neighbor: dp[i - 1][j]
2. The Left Neighbor: dp[i][j - 1]
3. The Diagonal Top-Left Neighbor: dp[i - 1][j - 1]

Real-Life Analogy: Building Blocks and Foundation Constraints

Imagine building a stone staircase or stacking blocks. To place a larger square block of size $k$ with its bottom-right corner at (i, j), the top, left, and diagonal supporting blocks must all be able to support a square of size at least $k-1$. If any of the three neighbors can only support a smaller square, the new square's growth is bottlenecked by that minimum neighbor.

The Recurrence Relation

If matrix[i][j] == 1:
   dp[i][j] = 1 + min(dp[i - 1][j], dp[i][j - 1], dp[i - 1][j - 1])
Else:
   dp[i][j] = 0

6. Complete Flow Walkthrough

Consider the initialization and population workflow for the matrix:

Base Cases (First Row & First Column)

• For any cell in the first row (i = 0), a square can only span of size 1 if the cell itself is 1. Thus, dp[0][j] = matrix[0][j].
• For any cell in the first column (j = 0), dp[i][0] = matrix[i][0].

7. Implementation Code (Java / Tabulation)

class Solution {
    public int countSquares(int[][] matrix) {
        int n = matrix.length;
        int m = matrix[0].length;
        int[][] dp = new int[n][m];
        int totalSquares = 0;

        // Step 1: Copy the first row and first column into the dp table
        for (int j = 0; j < m; j++) {
            dp[0][j] = matrix[0][j];
        }
        for (int i = 0; i < n; i++) {
            dp[i][0] = matrix[i][0];
        }

        // Step 2: Fill the rest of the DP table using the recurrence relation
        for (int i = 1; i < n; i++) {
            for (int j = 1; j < m; j++) {
                if (matrix[i][j] == 1) {
                    dp[i][j] = 1 + Math.min(dp[i - 1][j], 
                                  Math.min(dp[i][j - 1], dp[i - 1][j - 1]));
                } else {
                    dp[i][j] = 0;
                }
            }
        }

        // Step 3: Accumulate all values in the DP table to get total squares
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < m; j++) {
                totalSquares += dp[i][j];
            }
        }

        return totalSquares;
    }
}

8. Real-World Practical Pitfalls & Edge Cases

• Zero-Filled Matrices: If every element in the matrix is 0, the DP table will remain entirely filled with 0s, correctly returning a total sum of 0.
• Single Row or Single Column Matrices: If $n = 1$ or $m = 1$, the nested loop for inner cells is bypassed entirely, and the result is accurately computed just from the base row/column initialization copy.
• Overlooking Base Initialization: Forgetting to copy the first row and column prior to starting the inner loop at indices (1, 1) will result in ArrayIndexOutOfBoundsException or incorrect square aggregation.

9. Summary and Key Takeaways

• Problem Classification: Counting square submatrices with all ones on a 2D grid is an optimal candidate for 2D Dynamic Programming Tabulation.
• Definition of DP State: dp[i][j] denotes the maximum side length of a square whose bottom-right corner is positioned at (i, j).
• Total Summation: Summing every entry in the populated dp[][] matrix yields the exact total number of square submatrices of all possible sizes.
• Recurrence Formula: For any cell (i, j) containing 1, the value is calculated as 1 + min(top, left, diagonal).
• Base Cases: The first row and first column mirror the input matrix values directly since they cannot expand backward.
• Time Complexity: Evaluates in $O(n \times m)$ time by visiting each matrix cell a constant number of times.
• Space Complexity: Requires $O(n \times m)$ auxiliary space for the DP tracking table.