统计全为 1 的正方形子矩阵 - 动态规划解法

统计全为 1 的正方形子矩阵 - 动态规划解法

给你一个 m * n 的矩阵，矩阵中的元素不是 0 就是 1，请你统计并返回其中完全由 1 组成的正方形子矩阵的个数。

示例 1：

输入：matrix = {{0,1,1,1}, {1,1,1,1}, {0,1,1,1} }

输出：15

解释： 边长为 1 的正方形有 10 个。 边长为 2 的正方形有 4 个。 边长为 3 的正方形有 1 个。 正方形的总数 = 10 + 4 + 1 = 15.

示例 2：

输入：matrix = { {1,0,1}, {1,1,0}, {1,1,0}}

输出：7

算法：动态规划

设 dp[i][j] 表示以 (i,j) 为右下角的正方形的最大边长，则有：

• 若 matrix[i][j] == 0，则 dp[i][j] = 0
• 否则，dp[i][j] 的值取决于其左边，上边和左上角的三个位置的 dp 值，即 dp[i][j] = min(dp[i-1][j], dp[i][j-1], dp[i-1][j-1]) + 1

**时间复杂度：**O(mn)，其中 m 和 n 分别为矩阵的行数和列数。

**空间复杂度：**O(mn)，需要一个大小为 m*n 的二维数组存储 dp 值。

Python 代码：

def countSquares(matrix):
    m, n = len(matrix), len(matrix[0])
    dp = [[0 for _ in range(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

使用方法：

matrix = [[0, 1, 1, 1], [1, 1, 1, 1], [0, 1, 1, 1]]
print(countSquares(matrix))  # 输出 15