Medium

题目描述

给你一个 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
解释:
边长为 1 的正方形有 6 个。  
边长为 2 的正方形有 1 个。
正方形的总数 = 6 + 1 = 7.

提示:

  • 1 <= arr.length <= 300
  • 1 <= arr[0].length <= 300
  • 0 <= arr[i][j] <= 1

解题思路

这是一道典型的动态规划问题,关键是找到状态转移的规律。

核心思路: 对于矩阵中的每个位置 (i, j),我们定义 dp[i][j] 表示以 (i, j) 为右下角的正方形的最大边长。

状态转移方程:

  • 如果 matrix[i][j] == 0,则 dp[i][j] = 0
  • 如果 matrix[i][j] == 1,则 dp[i][j] = min(dp[i-1][j], dp[i][j-1], dp[i-1][j-1]) + 1

这个转移方程的含义是:要形成一个更大的正方形,当前位置必须为 1,同时它的左边、上边、左上角三个位置都能形成正方形,新正方形的边长就是这三个正方形边长的最小值加 1。

为什么这样可行? 每个位置的 dp 值不仅表示以该位置为右下角的最大正方形边长,同时也表示以该位置为右下角能形成多少个正方形。例如,如果 dp[i][j] = 3,说明以 (i,j) 为右下角可以形成边长为 1、2、3 的正方形各一个,总共 3 个。

最终答案就是所有位置的 dp 值之和。为了节省空间,我们可以直接在原矩阵上进行修改。

代码实现

class Solution {
public:
    int countSquares(vector<vector<int>>& matrix) {
        int m = matrix.size(), n = matrix[0].size();
        int result = 0;
        
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                if (matrix[i][j] == 1 && i > 0 && j > 0) {
                    matrix[i][j] = min({matrix[i-1][j], matrix[i][j-1], matrix[i-1][j-1]}) + 1;
                }
                result += matrix[i][j];
            }
        }
        
        return result;
    }
};
class Solution:
    def countSquares(self, matrix: List[List[int]]) -> int:
        m, n = len(matrix), len(matrix[0])
        result = 0
        
        for i in range(m):
            for j in range(n):
                if matrix[i][j] == 1 and i > 0 and j > 0:
                    matrix[i][j] = min(matrix[i-1][j], matrix[i][j-1], matrix[i-1][j-1]) + 1
                result += matrix[i][j]
        
        return result
public class Solution {
    public int CountSquares(int[][] matrix) {
        int m = matrix.Length, n = matrix[0].Length;
        int result = 0;
        
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                if (matrix[i][j] == 1 && i > 0 && j > 0) {
                    matrix[i][j] = Math.Min(Math.Min(matrix[i-1][j], matrix[i][j-1]), matrix[i-1][j-1]) + 1;
                }
                result += matrix[i][j];
            }
        }
        
        return result;
    }
}
/**
 * @param {number[][]} matrix
 * @return {number}
 */
var countSquares = function(matrix) {
    let m = matrix.length;
    let n = matrix[0].length;
    let count = 0;
    
    for (let i = 0; i < m; i++) {
        for (let j = 0; j < n; j++) {
            if (matrix[i][j] === 1) {
                if (i === 0 || j === 0) {
                    count += 1;
                } else {
                    matrix[i][j] = Math.min(matrix[i-1][j], matrix[i][j-1], matrix[i-1][j-1]) + 1;
                    count += matrix[i][j];
                }
            }
        }
    }
    
    return count;
};

复杂度分析

复杂度类型
时间复杂度O(m × n)
空间复杂度O(1)

其中 m 和 n 分别为矩阵的行数和列数。我们只需要遍历矩阵一次,每个位置的计算都是常数时间,所以时间复杂度为 O(m × n)。由于直接在原矩阵上修改,没有使用额外的存储空间,空间复杂度为 O(1)。

相关题目