[LeetCode 1277] Count Square Submatrices with All Ones

Given a m * n matrix of ones and zeros, return how many square submatrices have all ones.

Example 1:

Input: matrix =
[
  [0,1,1,1],
  [1,1,1,1],
  [0,1,1,1]
]
Output: 15
Explanation: 
There are 10 squares of side 1.
There are 4 squares of side 2.
There is  1 square of side 3.
Total number of squares = 10 + 4 + 1 = 15.

Example 2:

Input: matrix = 
[
  [1,0,1],
  [1,1,0],
  [1,1,0]
]
Output: 7
Explanation: 
There are 6 squares of side 1.  
There is 1 square of side 2. 
Total number of squares = 6 + 1 = 7.

Constraints:

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

The brute force solution of using each cell as the top left corner and check all possible square submatrices has runtime of O(N^3) * O(N^2), which is too slow.

O(N^2) dynamic programming solution

The final answer is the sum of the number of square submatrices at cell (i, j) over all cells. So let's define dp[i][j] as the maximum side length of the square submatrix whose bottom right corner is cell (i, j). The maximum side length also equals to the number of square submatrices with bottom right corner at cell (i, j). So dp[i][j] is also the number of square submatrices whose bottom right corner is cell (i, j).

If cell (i, j) is 0, then dp[i][j] is 0; Otherwise, to compute dp[i][j], we need to take the minimum side length of dp[i - 1][j], dp[i][j - 1] and dp[i - 1][j - 1] and add 1 to this min.

class Solution {
    public int countSquares(int[][] matrix) {
        int m = matrix.length, n = matrix[0].length, cnt = 0;  
        //dp[i][j]: the number of square submatrices whose bottom right corner is cell (i, j)
        //also the maximum side length of the square submatrix whose bottom right corner is cell (i, j)
        int[][] dp = new int[m][n];
        for(int i = 0; i < m; i++) {
            for(int j = 0; j < n; j++) {
                if(matrix[i][j] == 1 && i > 0 && j > 0) {
                    dp[i][j] = 1 + Math.min(Math.min(dp[i - 1][j], dp[i][j - 1]),dp[i - 1][j - 1]);
                }
                else {
                    dp[i][j] = matrix[i][j];
                }
                cnt += dp[i][j];
            }
        }
        return cnt;
    }
}