☕ DSA Java Articles

Topics › Dynamic Programming › DP on Squares

Count Square Submatrices with All Ones - (DP-56)

Easy TUF ID 2395 A2Z

LeetCode ↗ Java source ↗ Edit this article ↗

Problem

You are given an n × m matrix of 0s and 1s. Count how many square submatrices consist entirely of 1s. Squares of every size count, and overlapping squares are counted separately: a 2 × 2 block of ones contains four 1 × 1 squares and one 2 × 2 square, so it contributes 5.

Constraints (LeetCode 1277): 1 ≤ n, m ≤ 300, every entry is 0 or 1. The methods also return 0 for an empty matrix.

Example 1

0 1 1 1
1 1 1 1
0 1 1 1
answer = 15              // 10 squares of side 1, 4 of side 2, 1 of side 3

Example 2

1 0 1
1 1 0
1 1 0
answer = 7               // 6 squares of side 1, 1 of side 2 (bottom left)

Approach 1: Brute force

Intuition. Every square has a top-left corner and a side length. From each corner, grow the square one step at a time. When the side grows from s - 1 to s, the old square was already checked, so only the new bottom row and new right column need to be looked at. As soon as a 0 appears, stop: every larger square from the same corner contains that 0 too.

Algorithm.

  1. For every cell (i, j) as the top-left corner, for s = 1, 2, ... while the square fits:
  2. Check the s cells of row i + s - 1 and of column j + s - 1 inside the square.
  3. If all are 1, count the square and continue; otherwise break.
static int bruteForce(int[][] matrix) {
    int n = matrix.length, m = n == 0 ? 0 : matrix[0].length, total = 0;
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < m; j++) {
            for (int s = 1; i + s <= n && j + s <= m; s++) {
                boolean ok = true;                          // only the new bottom row and right column are unchecked
                for (int t = 0; t < s && ok; t++) {
                    if (matrix[i + s - 1][j + t] == 0 || matrix[i + t][j + s - 1] == 0) ok = false;
                }
                if (!ok) break;                             // a bigger square would contain this 0 too
                total++;
            }
        }
    }
    return total;
}

Complexity. Time O(n · m · L^2) with L = min(n, m): up to L sizes per corner, each step checking O(s) cells. On an all-ones 300 × 300 matrix that is more than a billion cell checks. Space O(1).

The key idea: count squares by their bottom-right corner

Every square has exactly one bottom-right corner. So if we can count, for each cell, how many all-ones squares end there, the answer is the sum of those counts.

Let dp[i][j] be the side of the largest all-ones square whose bottom-right corner is (i, j). If that side is s, then the squares of sides 1, 2, ..., s ending at (i, j) are all inside it and are all-ones, and no larger one exists. So exactly dp[i][j] squares end at (i, j), and

answer = sum of dp[i][j] over all cells

Now the recurrence for dp:

dp[i][j] = 0                                                    if matrix[i][j] == 0
dp[i][j] = 1                                                    if i == 0 or j == 0 (and the cell is 1)
dp[i][j] = 1 + min(dp[i-1][j], dp[i][j-1], dp[i-1][j-1])         otherwise

Why the minimum of those three neighbours? Look at a square of side s ending at (i, j). Remove its bottom row: what is left is a square of side s - 1 ending at (i - 1, j). Remove its right column instead: a square of side s - 1 ending at (i, j - 1). Remove both: one ending at (i - 1, j - 1). So all three neighbours are at least s - 1, which means s ≤ 1 + min(...).

The other direction: if all three neighbours are at least t, take the three t × t squares ending at them. The one ending above covers the top t rows of the right part, the one ending to the left covers the bottom t rows of the left part, and the diagonal one covers the top-left t × t block. Together with the cell (i, j) itself, they cover every cell of the (t + 1) × (t + 1) square ending at (i, j). So a square of side 1 + min(...) really exists. The weakest neighbour is the bottleneck.

Approach 2: Recursion

Intuition. Compute dp[i][j] directly with the recurrence as a recursive function solve(i, j), and add up solve(i, j) for every cell.

Algorithm.

  1. If the cell is 0, return 0. If it is on the first row or column, return 1.
  2. Otherwise return 1 + min of the three recursive neighbour calls.
  3. Sum the results over all cells.
static int recursive(int[][] matrix) {
    int total = 0;
    for (int i = 0; i < matrix.length; i++)
        for (int j = 0; j < matrix[0].length; j++) total += solve(i, j, matrix);
    return total;
}

static int solve(int i, int j, int[][] matrix) {
    if (matrix[i][j] == 0) return 0;
    if (i == 0 || j == 0) return 1;                         // first row or column: only the 1x1 square fits
    return 1 + Math.min(solve(i - 1, j, matrix), Math.min(solve(i, j - 1, matrix), solve(i - 1, j - 1, matrix)));
}

Complexity. Time exponential: each call makes three more and the same cell is recomputed along many paths. In an all-ones matrix the number of calls for cell (i, j) grows exponentially in i + j (it is bounded by 3^(i + j)). Space O(n + m) for the recursion depth.

Approach 3: Memoization

Intuition. There are only n · m distinct cells, so store each solve(i, j) in dp[i][j] the first time it is computed, with -1 as “not computed”. The cheap base cases (a 0 cell, or the first row or column) are answered before the cache is consulted.

Algorithm.

  1. Fill dp with -1.
  2. Same recursion, returning the cached value when present.
  3. Sum over all cells.
static int memoization(int[][] matrix) {
    int n = matrix.length, m = n == 0 ? 0 : matrix[0].length, total = 0;
    int[][] dp = new int[n][m];
    for (int[] row : dp) Arrays.fill(row, -1);
    for (int i = 0; i < n; i++)
        for (int j = 0; j < m; j++) total += memo(i, j, matrix, dp);
    return total;
}

static int memo(int i, int j, int[][] matrix, int[][] dp) {
    if (matrix[i][j] == 0) return 0;
    if (i == 0 || j == 0) return 1;
    if (dp[i][j] != -1) return dp[i][j];
    return dp[i][j] = 1 + Math.min(memo(i - 1, j, matrix, dp), Math.min(memo(i, j - 1, matrix, dp), memo(i - 1, j - 1, matrix, dp)));
}

Complexity. Time O(nm): each cell is computed once in O(1), and the outer loop adds O(nm). Space O(nm) for the table plus up to O(n + m) stack.

Approach 4: Tabulation

Intuition. dp[i][j] reads the cell above, the cell to the left, and the cell diagonally up-left. Row-major order, top to bottom and left to right, visits all three before (i, j). The sum can be accumulated in the same loop.

Algorithm.

  1. For each cell in row-major order, apply the three-case recurrence.
  2. Add dp[i][j] to the total.
static int tabulation(int[][] matrix) {
    int n = matrix.length, m = n == 0 ? 0 : matrix[0].length, total = 0;
    int[][] dp = new int[n][m];
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < m; j++) {
            if (matrix[i][j] == 0) dp[i][j] = 0;
            else if (i == 0 || j == 0) dp[i][j] = 1;
            else dp[i][j] = 1 + Math.min(dp[i - 1][j], Math.min(dp[i][j - 1], dp[i - 1][j - 1]));
            total += dp[i][j];                              // dp[i][j] squares have (i, j) as bottom-right corner
        }
    }
    return total;
}

Dry run on Example 1. The table dp is:

0 1 1 1          row sum 3
1 1 2 2          row sum 6
0 1 2 3          row sum 6

For example, dp[1][2] = 1 + min(dp[0][2], dp[1][1], dp[0][1]) = 1 + min(1, 1, 1) = 2, and dp[2][3] = 1 + min(dp[1][3], dp[2][2], dp[1][2]) = 1 + min(2, 2, 2) = 3, the single 3 × 3 square. The total is 3 + 6 + 6 = 15. Grouped by size: 10 cells have dp ≥ 1, 4 have dp ≥ 2, and 1 has dp ≥ 3, giving the 10 + 4 + 1 from the example.

Complexity. Time O(nm), space O(nm).

Approach 5: Space optimised, one row

Intuition. Each row reads only the previous row and itself. Keep a single array where, while row i is being filled, the cells to the left of j already hold row i and the cells from j onward still hold row i - 1. Before overwriting position j, its old value is the “above” neighbour; the “diagonal” neighbour is the old value of position j - 1, which we saved one step earlier in diag.

The array has m + 1 cells and is shifted by one: dp[j + 1] stores column j, and dp[0] stays 0 forever. That extra cell acts as a wall, so the first-column case needs no special code: 1 + min(above, 0, 0) = 1 for a 1. The first row works the same way, because the array starts as zeros.

Algorithm.

  1. For each row, reset diag = 0.
  2. For each j: save above = dp[j + 1]; set dp[j + 1] = 0 for a 0 cell, else 1 + min(above, dp[j], diag); then diag = above; add dp[j + 1] to the total.
static int spaceOptimized(int[][] matrix) {
    int n = matrix.length, m = n == 0 ? 0 : matrix[0].length, total = 0;
    int[] dp = new int[m + 1];                              // dp[j + 1] is column j; dp[0] stays 0 as a wall
    for (int i = 0; i < n; i++) {
        int diag = 0;                                       // previous row's value at column j - 1
        for (int j = 0; j < m; j++) {
            int above = dp[j + 1];                          // still the previous row's value
            dp[j + 1] = matrix[i][j] == 0 ? 0 : 1 + Math.min(above, Math.min(dp[j], diag));
            diag = above;
            total += dp[j + 1];
        }
    }
    return total;
}

Complexity. Time O(nm), space O(m). If modifying the input is allowed, the recurrence can even be written into matrix itself, using O(1) extra space.

Edge cases and pitfalls