leetcode 1277 统计全为1的正方形子矩阵（动态规划）

https://leetcode-cn.com/problems/count-square-submatrices-with-all-ones/

题目
给你一个 m * n 的矩阵，矩阵中的元素不是 0 就是 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.

输入：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

思路
动态规划是真的难啊。
用 dp[m][n] 数组来统计，dp[i][j] 表示以 (i, j) 为右下角往左上延伸的最大的全为1的正方形的边长，然后全部累加起来就是总数，因为右下角是确定的，并且边长为 k，即该元素值为 k 时，代表一定有以该点为右下角的边长为 1，2 . . . k 等 k 个正方形，边长和一个顶点确定，不会造成统计的重复。
状态转移方程为 dp[i][j] = matrix[i][j] ? min(min(dp[i - 1][j - 1], dp[i][j - 1]), dp[i - 1][j]) + 1 : 0
代码

#include <cmath>
class Solution {
public:
    int countSquares(vector<vector<int>>& matrix) {
        int ans = 0;
        int m = matrix.size();
        int n = matrix[0].size();
        vector <vector<int>> dp(m, vector<int> (n, 0));
        for (int i = 0; i < m; i++) {
            dp[i][0] = matrix[i][0];
            ans += matrix[i][0];
        }
        for (int i = 0; i < n; i++) {
            dp[0][i] = matrix[0][i];
            ans += matrix[0][i];
        } 
        ans -= matrix[0][0];
        for (int i = 1; i < m; i++) 
            for (int j = 1; j < n; j++) {
                if (!matrix[i][j]) continue;
                dp[i][j] = 1;
                dp[i][j] += min(min(dp[i - 1][j - 1],dp[i][j - 1]), dp[i - 1][j]);
                ans += dp[i][j];
            }
        return ans;
    }
};