Program to count number of square submatrices with all ones using Python

In this problem, we need to count all possible square submatrices that contain only ones in a binary matrix. This is solved using dynamic programming where we track the largest square ending at each position.

Problem Understanding

Given a binary matrix, we count squares of all sizes. For the example matrix ?

Algorithm Approach

We use dynamic programming where each cell stores the size of the largest square ending at that position. The key insight is ?

• For the first row and column, a cell can only form a 1×1 square if it contains 1

• For other cells, the largest square size is determined by the minimum of three neighbors plus 1

Implementation

def count_square_submatrices(matrix):
    if not matrix or not matrix[0]:
        return 0
    
    rows = len(matrix)
    cols = len(matrix[0])
    result = 0
    
    for row in range(rows):
        for col in range(cols):
            if row == 0 or col == 0:
                # First row or column - can only form 1x1 squares
                if matrix[row][col] == 1:
                    result += 1
            elif matrix[row][col] == 1:
                # Find minimum of three neighbors and add 1
                square_size = min(
                    matrix[row-1][col],      # top
                    matrix[row][col-1],      # left  
                    matrix[row-1][col-1]     # diagonal
                ) + 1
                
                matrix[row][col] = square_size
                result += square_size
    
    return result

# Test with the example
matrix = [[0,1,1,1],[1,1,1,1],[0,1,1,1]]
print(f"Number of square submatrices: {count_square_submatrices(matrix)}")

Number of square submatrices: 15

How It Works

The algorithm transforms the matrix to store the maximum square size ending at each position ?

def count_with_visualization(matrix):
    rows, cols = len(matrix), len(matrix[0])
    result = 0
    
    print("Original matrix:")
    for row in matrix:
        print(row)
    
    print("\nAfter processing (shows max square size at each position):")
    
    for row in range(rows):
        for col in range(cols):
            if row == 0 or col == 0:
                if matrix[row][col] == 1:
                    result += 1
            elif matrix[row][col] == 1:
                square_size = min(
                    matrix[row-1][col],
                    matrix[row][col-1], 
                    matrix[row-1][col-1]
                ) + 1
                matrix[row][col] = square_size
                result += square_size
    
    for row in matrix:
        print(row)
    
    print(f"\nTotal squares: {result}")
    return result

# Test
matrix = [[0,1,1,1],[1,1,1,1],[0,1,1,1]]
count_with_visualization(matrix)

Original matrix:
[0, 1, 1, 1]
[1, 1, 1, 1]
[0, 1, 1, 1]

After processing (shows max square size at each position):
[0, 1, 1, 1]
[1, 1, 2, 2]
[0, 1, 2, 3]

Total squares: 15

Time and Space Complexity

Conclusion

This dynamic programming solution efficiently counts all square submatrices by storing the maximum square size at each position. Each cell's value represents how many squares end at that position, and the sum gives the total count.