Problem Description

You are given a m * n matrix filled with ones (1) and zeros (0). Your task is to count how many square submatrices contain all ones.

A square submatrix is a submatrix where the number of rows equals the number of columns. For example, a 1×1, 2×2, or 3×3 submatrix would all be square submatrices.

The problem asks you to find the total count of all possible square submatrices of any size that contain only 1s.

For instance, if you have a 3×3 matrix:

1 0 1
1 1 0
1 1 0

The square submatrices with all ones would be:

• Four 1×1 squares at positions (0,0), (1,0), (1,1), (2,0), (2,1)
• One 2×2 square with top-left corner at (1,0)

So the total count would be 6.

The solution uses dynamic programming where f[i][j] represents the side length of the largest square submatrix with all ones that has its bottom-right corner at position (i,j). The key insight is that if matrix[i][j] = 1, then f[i][j] depends on the minimum of three adjacent positions (top, left, and top-left diagonal) plus 1. This is because a square can only extend as far as the smallest square that can be formed from these three positions. The sum of all f[i][j] values gives the total count of square submatrices, as each f[i][j] contributes all squares of size 1×1 up to f[i][j]×f[i][j] that end at position (i,j).

Intuition

The key observation is that every square submatrix has a bottom-right corner. If we can figure out, for each position (i,j), how many square submatrices have their bottom-right corner at that position, we can sum them all up to get our answer.

Let's think about what determines the largest square that can have its bottom-right corner at position (i,j). If matrix[i][j] = 0, clearly no square of ones can end there. But if matrix[i][j] = 1, we need to look at what's happening around it.

For a 2×2 square to exist with bottom-right at (i,j), we need:

• The current cell (i,j) to be 1
• The cell above (i-1,j) to be 1
• The cell to the left (i,j-1) to be 1
• The diagonal cell (i-1,j-1) to be 1

For a 3×3 square, we need all cells in a 3×3 region to be 1. But here's the clever insight: instead of checking all 9 cells, we can build upon smaller squares. A 3×3 square with bottom-right at (i,j) can exist only if:

• There's at least a 2×2 square ending at (i-1,j) (above)
• There's at least a 2×2 square ending at (i,j-1) (left)
• There's at least a 2×2 square ending at (i-1,j-1) (diagonal)

The maximum square size at (i,j) is limited by the smallest of these three neighboring squares, plus 1 for the current cell. Why the minimum? Because if any of these three positions can only support a smaller square, it creates a "bottleneck" that prevents us from forming a larger square at (i,j).

This leads us to the recurrence relation: f[i][j] = min(f[i-1][j-1], f[i-1][j], f[i][j-1]) + 1 when matrix[i][j] = 1.

The beauty of this approach is that f[i][j] not only tells us the side length of the largest square ending at (i,j), but also implicitly counts all smaller squares. If f[i][j] = 3, it means there's a 3×3 square, a 2×2 square, and a 1×1 square all ending at position (i,j). So we add f[i][j] to our total count.

Learn more about Dynamic Programming patterns.

Solution Approach

We implement the solution using dynamic programming with a 2D table f where f[i][j] represents the side length of the largest square submatrix with all ones that has its bottom-right corner at position (i,j).

Initialization:

• Create a 2D array f with dimensions m × n initialized with zeros
• Initialize a counter ans to track the total count of square submatrices

State Transition: We iterate through each cell (i,j) in the matrix:

1. When matrix[i][j] = 0:

   • Skip this cell (or keep f[i][j] = 0) since no square of ones can end here

2. When matrix[i][j] = 1:

   • Base case (first row or first column): If i = 0 or j = 0, then f[i][j] = 1. This is because cells in the first row or column can only form 1×1 squares as they don't have enough neighbors above or to the left.

   • General case: For all other positions, we apply the recurrence relation:

     f[i][j] = min(f[i-1][j-1], f[i-1][j], f[i][j-1]) + 1

     This formula checks:

     • f[i-1][j-1]: largest square ending at the diagonal position
     • f[i-1][j]: largest square ending at the position above
     • f[i][j-1]: largest square ending at the position to the left

     We take the minimum because the square at (i,j) can only extend as far as the smallest neighboring square allows.

3. Accumulate the result:

   • After computing f[i][j], add it to ans. This works because if f[i][j] = k, it means there are exactly k square submatrices ending at position (i,j): one of size 1×1, one of size 2×2, ..., and one of size k×k.

Time Complexity: O(m × n) where we visit each cell once
Space Complexity: O(m × n) for the DP table

The algorithm efficiently counts all square submatrices by building up from smaller squares to larger ones, avoiding redundant calculations through dynamic programming.

Example Walkthrough

Let's walk through a small example with this matrix:

1 1 1
1 1 1
1 1 0

We'll build our DP table f[i][j] step by step, where f[i][j] represents the side length of the largest square ending at position (i,j).

Step 1: Initialize and process first row (i=0)

• f[0][0] = 1 (matrix[0][0] = 1, first row so can only be 1×1)
• f[0][1] = 1 (matrix[0][1] = 1, first row so can only be 1×1)
• f[0][2] = 1 (matrix[0][2] = 1, first row so can only be 1×1)
• Running total: ans = 1 + 1 + 1 = 3

Step 2: Process second row (i=1)

• f[1][0] = 1 (matrix[1][0] = 1, first column so can only be 1×1)
• f[1][1] = min(f[0][0], f[0][1], f[1][0]) + 1 = min(1, 1, 1) + 1 = 2
  • This means a 2×2 square ends here, contributing 2 squares (one 1×1 and one 2×2)
• f[1][2] = min(f[0][1], f[0][2], f[1][1]) + 1 = min(1, 1, 2) + 1 = 2
  • Another 2×2 square ends here
• Running total: ans = 3 + 1 + 2 + 2 = 8

Step 3: Process third row (i=2)

• f[2][0] = 1 (matrix[2][0] = 1, first column so can only be 1×1)
• f[2][1] = min(f[1][0], f[1][1], f[2][0]) + 1 = min(1, 2, 1) + 1 = 2
  • A 2×2 square ends here
• f[2][2] = 0 (matrix[2][2] = 0, so no squares can end here)
• Running total: ans = 8 + 1 + 2 + 0 = 11

Final DP table:

f = [1, 1, 1]
    [1, 2, 2]
    [1, 2, 0]

Answer: 11 square submatrices

Breaking down what we counted:

• Nine 1×1 squares (every cell with value 1)
• Two 2×2 squares ending at (1,1) and (1,2)
• No 3×3 squares (blocked by the 0 at position (2,2))

The key insight is that each f[i][j] value tells us how many squares end at that position. When f[1][1] = 2, it means both a 1×1 and a 2×2 square end at (1,1), so we add 2 to our total count.

Solution Implementation

Time and Space Complexity

Time Complexity: O(m × n)

The algorithm iterates through each element of the matrix exactly once using nested loops. The outer loop runs m times (number of rows), and for each iteration, the inner loop runs n times (number of columns). Within each iteration, all operations performed are constant time O(1) - including the comparison checks, minimum calculation among three values, and arithmetic operations. Therefore, the overall time complexity is O(m × n).

Space Complexity: O(m × n)

The algorithm creates a 2D array f with dimensions m × n to store the dynamic programming values, where f[i][j] represents the side length of the largest square with bottom-right corner at position (i, j). This auxiliary array requires m × n space. The additional variables (m, n, ans, i, j, v) use constant space O(1). Therefore, the overall space complexity is O(m × n).

Learn more about how to find time and space complexity quickly.

Common Pitfalls

1. Incorrect Initialization for First Row and Column

A common mistake is forgetting to handle the first row and column as special cases. Some developers might incorrectly try to access dp[i-1][j] or dp[i][j-1] when i=0 or j=0, leading to index out of bounds errors or incorrect logic.

Incorrect approach:

# Wrong: Doesn't handle edge cases
for i in range(rows):
    for j in range(cols):
        if matrix[i][j] == 1:
            dp[i][j] = min(dp[i-1][j-1], dp[i-1][j], dp[i][j-1]) + 1
            # This fails when i=0 or j=0!

Correct approach:

# Right: Explicitly handles first row and column
if matrix[i][j] == 1:
    if i == 0 or j == 0:
        dp[i][j] = 1
    else:
        dp[i][j] = min(dp[i-1][j-1], dp[i-1][j], dp[i][j-1]) + 1

2. Space Optimization Confusion

When attempting to optimize space from O(m×n) to O(n) by using only one or two rows, developers often mix up the previous row's values with the current row's values, leading to incorrect calculations.

Problematic space-optimized version:

# Wrong: Overwrites values needed for future calculations
dp = [0] * cols
for i in range(rows):
    for j in range(cols):
        if matrix[i][j] == 1:
            if i == 0 or j == 0:
                dp[j] = 1
            else:
                # dp[j-1] is from current row, dp[j] is from previous row
                # but we need diagonal value which gets lost!
                dp[j] = min(???, dp[j], dp[j-1]) + 1

Correct space-optimized approach:

# Right: Keep track of diagonal value separately
dp = [0] * cols
for i in range(rows):
    prev_diagonal = 0
    for j in range(cols):
        temp = dp[j]  # Save current value before updating
        if matrix[i][j] == 1:
            if i == 0 or j == 0:
                dp[j] = 1
            else:
                dp[j] = min(prev_diagonal, dp[j], dp[j-1]) + 1
        else:
            dp[j] = 0
        prev_diagonal = temp  # Update diagonal for next iteration
        total_squares += dp[j]

3. Misunderstanding the DP Value Meaning

Some developers mistakenly think dp[i][j] represents the total count of squares ending at (i,j) rather than the side length of the largest square. This leads to incorrect accumulation logic.

Wrong interpretation:

# Incorrect: Treating dp[i][j] as count instead of side length
dp[i][j] = dp[i-1][j-1] + dp[i-1][j] + dp[i][j-1] + 1  # Wrong formula!
total_squares = dp[rows-1][cols-1]  # Wrong accumulation!

Correct understanding:

# Right: dp[i][j] is the side length, not the count
dp[i][j] = min(dp[i-1][j-1], dp[i-1][j], dp[i][j-1]) + 1
total_squares += dp[i][j]  # Add side length (which equals count of squares ending here)

4. Forgetting to Handle Zero Values

A subtle pitfall is not properly resetting dp[i][j] to 0 when matrix[i][j] = 0. While the provided solution skips these cells (which implicitly keeps them at 0), explicitly setting them can prevent bugs when modifying the code.

More explicit and safer:

if matrix[i][j] == 0:
    dp[i][j] = 0  # Explicitly set to 0 for clarity
    continue

These pitfalls highlight the importance of understanding the DP state definition, properly handling boundary conditions, and being careful when optimizing for space complexity.