Count Square Submatrices with All Ones

1public int countSquares(int[][] matrix) {
2    int m = matrix.length;
3    int n = matrix[0].length;
4    int[][] dp = new int[m][n];
5    int count = 0;
6
7    for (int i = 0; i < m; i++) {
8        for (int j = 0; j < n; j++) {
9            // Base case: First row or first column
10            if (i == 0 || j == 0) {
11                dp[i][j] = matrix[i][j];
12            } else if (matrix[i][j] == 1) {
13                // Determine min square size possible
14                dp[i][j] = Math.min(dp[i - 1][j], Math.min(dp[i][j - 1], dp[i - 1][j - 1])) + 1;
15            }
16            // Add to total count
17            count += dp[i][j];
18        }
19    }
20
21    return count;
22}