Count Square Submatrices with All Ones

medium dp matrix

Problem

Given a 0/1 matrix, count how many square submatrices contain only 1s. dp[i][j] equals the side length of the largest such square whose bottom-right corner is (i,j); summing all dp values gives the total count.

Inputmatrix = [[0,1,1,1],[1,1,1,1],[0,1,1,1]]

Output15

10 squares of side 1, 4 of side 2, and 1 of side 3.

matrix (rows separated by ';' e.g. 0,1,1,1;1,1,1,1;0,1,1,1)

def count_squares(matrix):
    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] = 1 + min(dp[i-1][j-1], dp[i-1][j], dp[i][j-1])
                total += dp[i][j]
    return total

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

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

int countSquares(vector<vector<int>>& matrix) {
    int m = matrix.size(), n = matrix[0].size();
    vector<vector<int>> dp(m, vector<int>(n, 0));
    int total = 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] = 1 + min({ dp[i-1][j-1], dp[i-1][j], dp[i][j-1] });
                total += dp[i][j];
            }
        }
    }
    return total;
}

Explanation

The clever counting trick: dp[i][j] = the side length of the largest all-1 square whose bottom-right corner sits at (i, j). A cell with dp[i][j] = 3 is the corner of one 3x3, one 2x2, and one 1x1 square — so it contributes exactly 3 squares. That means the answer is just the sum of every dp value.

If matrix[i][j] is 0, no square ends there, so dp[i][j] stays 0. Cells on the top row or left column can only ever form a 1x1, so they get dp = 1 when they hold a 1.

For an inner cell holding a 1, the square is limited by its three neighbours: dp[i][j] = 1 + min(dp[i-1][j-1], dp[i-1][j], dp[i][j-1]). The smallest neighbour is the bottleneck — a bigger square needs the top, left, and diagonal squares to all be at least that big.

As we sweep the grid, we add each dp[i][j] into a running total, which gives the count of all all-1 squares.

Example: matrix = [[0,1,1,1],[1,1,1,1],[0,1,1,1]]. The dp values sum to 15 — that's 10 squares of side 1, 4 of side 2, and 1 of side 3.

Time: O(m·n) Space: O(m·n)