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 <= 3001 <= arr[0].length <= 3000 <= arr[i][j] <= 1class Solution {
private void print(int[][] mat) {
for (int i = 0; i < mat.length; i++)
System.out.println(Arrays.toString(mat[i]));
System.out.println();
}
public int countSquares(int[][] matrix) {
int m = matrix.length, n = matrix[0].length, maxDim = Math.min(m, n), count = 0;
// generate prefix sum matrix
int[][] prefix = new int[m + 1][n + 1];
for (int i = 0; i < m; i++) {
int sum = 0;
for (int j = 0; j < n; j++)
prefix[i + 1][j + 1] = sum += matrix[i][j];
}
// print(prefix);
for (int j = 0; j <= n; j++)
for (int i = 0; i < m; i++)
prefix[i + 1][j] += prefix[i][j];
// print(prefix);
for (int k = 1; k <= maxDim; k++) {
// System.out.println("k->" + k);
for (int i = k; i <= m; i++) {
for (int j = k; j <= n; j++) {
// System.out.print(prefix[i][j] - prefix[i - k][j] - prefix[i][j - k] + prefix[i - k][j - k] + "\t");
if (prefix[i][j] - prefix[i - k][j] - prefix[i][j - k] + prefix[i - k][j - k] == k * k)
count++;
}
// System.out.println();
}
}
return count;
}
}