Code Implementation

class Solution:
    def countSquares(self, matrix):
        if not matrix or not matrix[0]:
            return 0
        m, n = len(matrix), len(matrix[0])
        dp = [[0]*n for _ in range(m)]
        total = 0
        for i in range(m):
            for j in range(n):
                if matrix[i][j] == 1:
                    if i == 0 or j == 0:
                        dp[i][j] = 1
                    else:
                        dp[i][j] = min(dp[i-1][j], dp[i][j-1], dp[i-1][j-1]) + 1
                    total += dp[i][j]
        return total
      

class Solution {
public:
    int countSquares(vector<vector<int>>& matrix) {
        int m = matrix.size(), n = matrix[0].size(), total = 0;
        vector<vector<int>> dp(m, vector<int>(n, 0));
        for (int i = 0; i < m; ++i) {
            for (int j = 0; j < n; ++j) {
                if (matrix[i][j] == 1) {
                    if (i == 0 || j == 0)
                        dp[i][j] = 1;
                    else
                        dp[i][j] = min({dp[i-1][j], dp[i][j-1], dp[i-1][j-1]}) + 1;
                    total += dp[i][j];
                }
            }
        }
        return total;
    }
};
      

class Solution {
    public int countSquares(int[][] matrix) {
        int m = matrix.length, n = matrix[0].length, total = 0;
        int[][] dp = new int[m][n];
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                if (matrix[i][j] == 1) {
                    if (i == 0 || j == 0)
                        dp[i][j] = 1;
                    else
                        dp[i][j] = Math.min(Math.min(dp[i-1][j], dp[i][j-1]), dp[i-1][j-1]) + 1;
                    total += dp[i][j];
                }
            }
        }
        return total;
    }
}
      

var countSquares = function(matrix) {
    let m = matrix.length, n = matrix[0].length, total = 0;
    let dp = Array.from({length: m}, () => Array(n).fill(0));
    for (let i = 0; i < m; i++) {
        for (let j = 0; j < n; j++) {
            if (matrix[i][j] === 1) {
                if (i === 0 || j === 0)
                    dp[i][j] = 1;
                else
                    dp[i][j] = Math.min(dp[i-1][j], dp[i][j-1], dp[i-1][j-1]) + 1;
                total += dp[i][j];
            }
        }
    }
    return total;
};
      

Problem Description

You are given a 2D binary matrix called matrix with dimensions m x n. Each cell contains either a 0 or a 1. The task is to count the total number of square submatrices (of any size) that are filled entirely with 1s.

A square submatrix is a square region within the matrix, aligned with the axes, where every cell is 1. The squares can be of any size (1x1, 2x2, ..., up to min(m, n)xmin(m, n)), and can overlap.

Constraints:

• Each cell can be used in multiple squares (overlap is allowed).
• Count all possible squares, not just the largest ones.
• Matrix size: up to 300 x 300.

Thought Process

The most straightforward way to solve this problem is to check every possible square submatrix and verify if all its cells are 1. However, this brute-force approach would be very slow for large matrices, since there are many possible squares to check.

To optimize, we should look for ways to avoid redundant checks. Notice that if a cell is the bottom-right corner of a square of size k, it must also be the bottom-right corner of squares of size k-1, k-2, ..., 1. This suggests that we can build up the solution using smaller squares to help find larger ones.

Dynamic programming is a good fit for this type of "build up from smaller subproblems" scenario.

Solution Approach

We use a dynamic programming (DP) approach to efficiently count all square submatrices with all ones:

• Step 1: DP Table Definition
  Create a DP table dp of the same size as the input matrix. Each cell dp[i][j] represents the size of the largest square that ends at cell (i, j) (i.e., with (i, j) as its bottom-right corner), and all its cells are 1.
• Step 2: Transition
  For each cell (i, j) in the matrix:
  • If matrix[i][j] == 0, then dp[i][j] = 0 (no square ends here).
  • If matrix[i][j] == 1 and i == 0 or j == 0, then dp[i][j] = 1 (only a 1x1 square fits).
  • Otherwise, dp[i][j] = min(dp[i-1][j], dp[i][j-1], dp[i-1][j-1]) + 1. This is because the cell can extend the minimal square among its top, left, and top-left neighbors.
• Step 3: Counting
  For every cell, the value in dp[i][j] tells you how many squares (of sizes 1 up to dp[i][j]) end at that cell. Add up all dp[i][j] values to get the total number of square submatrices with all ones.

This approach ensures that each cell is processed only once, and we avoid redundant checks for overlapping squares.

Example Walkthrough

Let's walk through the following example:

    Input matrix:
    1 0 1
    1 1 0
    1 1 0
  

• Initialize a DP table with the same shape, filled with zeros.
• Process each cell:
  • (0,0): matrix=1 ⇒ dp=1
  • (0,1): matrix=0 ⇒ dp=0
  • (0,2): matrix=1 ⇒ dp=1
  • (1,0): matrix=1 ⇒ dp=1
  • (1,1): matrix=1 ⇒ min(1,0,0)+1=1 ⇒ dp=1
  • (1,2): matrix=0 ⇒ dp=0
  • (2,0): matrix=1 ⇒ dp=1
  • (2,1): matrix=1 ⇒ min(1,1,1)+1=2 ⇒ dp=2
  • (2,2): matrix=0 ⇒ dp=0
• Final DP table:

          1 0 1
          1 1 0
          1 2 0
        

• Sum of all DP values: 1+0+1+1+1+0+1+2+0 = 7.

Thus, there are 7 square submatrices with all ones in this example.

Time and Space Complexity

• Brute-force approach: For every possible square (of every possible size and position), check if all cells are 1. There are O(mn) starting points and up to O(min(m, n)) sizes per point, so total work is O(mn * min(m, n)^2).
• Optimized DP approach: We process each cell once, and each DP update is O(1). So, total time is O(mn).
• Space Complexity: The DP table uses O(mn) space. If we modify the input matrix in place (allowed in some scenarios), we can reduce to O(1) extra space.

Summary

By leveraging dynamic programming, we efficiently count all square submatrices with all ones by building up the solution from smaller squares. Instead of checking every possible square, we use the insight that the largest square ending at a cell depends on its neighbors. This reduces the time complexity from potentially cubic to quadratic, making the solution practical for large matrices. The elegance comes from recognizing overlapping subproblems and using DP to avoid redundant work.