Count Square Submatrices with all Ones

A brute-force approach is to iterate over all the square submatrices and check if they contain all ones.

This approach is too slow. The 2nd part of determining if this square submatrix contains all ones can be optimized.  We can store prefix sums for each cell. Then we can check if this submatrix has all ones in constant time. 

How to calculate prefix sums for a matrix?

Suppose we are in a cell (r, c). Then we can add prefix sums of (r - 1, c) and (r, c - 1) along with arr[r][c], but we have counted (r - 1, c - 1) two times. So we will subtract that once.

So, the relation becomes pref[r][c] = arr[r][c] + pref[r - 1][c] + pref[r][c - 1] - pref[r - 1][c - 1] (where pref[i][j] denotes the sum of the submatrix which has the top left corner at (0, 0) and the bottom right corner at (i, j)

Note that, for the first row, or first column, some of the terms do not exist. So, we can take them 0.

The steps are as follows:

1. Calculate the prefix sum matrix for the given matrix. Also, declare and initialize a variable ‘ans’ to ‘0’.
2. Iterate over all the upper-left corners of the square. Then iterate over the size of the square to check if it contains all ones.
3. Suppose the current upper-left corner has coordinates i, j, and the size of the square is s * s. Then the bottom-right corner has coordinates will be i + s - 1, j + s - 1.
4. Find the sum of this square using prefix sums and check if this sum equals s * s. If it is, then increment the answer count.
5. Return the answer.

First of all, observe that, for each cell, if we find the largest-sized square ending at that cell, then we can add this directly to the answer count because a larger sized-submatrix is not possible, and all the square submatrices from size 1 to MAX_SIZE( min( N, M ) ) is possible.

How to find it, though?

If the value of the current cell is 0, its contribution is zero since it cannot be the bottom right corner of any square.

If the current cell value is 1, and its coordinates are (r, c), then suppose the size of the maximum-sized square submatrix ending at the cell is S, then think about what should the values of the cells (r - 1, c), (r, c - 1), (r - 1, c - 1) be?

They all should be ≥ S - 1 because if any of them is smaller than S - 1, then at least one of the cells in this S * S submatrix will have 0.

For example, consider that we argue that (r, c) has a maximum-sized square value of 4. And let’s say, (r - 1, c) has value 2, (r, c - 1) has value 3, (r - 1, c - 1) also has value 3. What does this mean?

The above table lists the values, which must be 1. Now, if (r, c) has a maximum-sized square value of 4, this will force the only left out cell value to be 1. But then we got a contradiction: (r - 1, c) has value 2, but it should have value 3.

Convince yourself that the above fact is true.

That means the value for the cell (r, c) depends on whichever of the above three is minimum by one less.

The steps are as follows:

1. Declare and initialize a variable ‘ans’ to 0.
2. Iterate over all the bottom-right corners of the matrix. Two cases arise:

• If the current cell value is 0, then continue.
• Otherwise, we can take the minimum of three values as described above and add 1 to it.
• DP[i][j] = min({DP[i - 1][j], DP[i - 1][j - 1], DP[i][j - 1]) + 1.

3. Add DP[i][j] to the ans and return.