[Java] Find all square submatrices of sum K or greater

Here is an interesting problem, given a non necessarily square matrix and a threshold value K, find all square submatrices whose sum is greater than or equal to K.

Example:

111

111

111

If K = 1, we have:

1x1 submatrices = 9 of value 1

2x2 submatrices = 4 of value 4

3x3 submatrices = 1 of value 9

total is 14 submatrices whose sum is more or equal to K. 

One approach is to consider each cell as the top left corner of a square submatrix. From there we can then expand to squares of increasing sizes while counting the total sum, until the square would go out of bounds. We can do this in O((NM)^2) time as for each of the NM cells we do NM work.

We can notice that we do not need to recalculate the whole square when we increase it since the bigger square includes the smaller one, we can therefore only do the work necessary to include the new values, we cover with the new size:

11#

11#

###

To compute from the 2x2 square the bigger 3x3, we only need to add the cells marked with #.

CAUTION: notice the bottom right would be added twice, therefore we sum all # cells and subtract the value of bottom right once.

This is done with two separate loops both of which have O(min(N,M)) time for up to min(N,M) times, giving us O(min(N,M)^2) time for a total of O(NMmin(N,M))^2) time and O(1) space.

If we precompute the prefix matrix in O(NM) time using extra O(NM) space, we can reuse already calculated submatrices instead. In this case, we consider each cell as the bottom right cell of a square submatrix. For each of the NM cells, we try to expand from there using data from the prefix matrix to quickly calculate the value of each square. This calculation costs O(min(N,M)) and in total we repeat it at most min(N,M) times.

The resulting time then becomes O(NM + min(N,M)^2) at the cost of extra O(NM) space.

Maybe there is some clever mathematical observation to be done using additional space/structures that could bring the time down to O(NM) but I couldn't figure it out yet.

You can check my implementation of numSquaresOfSumK and numSquaresOfSumKDP with the helper function prefixSum to calculate the prefix matrix on my Gist along with some tests in NumSquaresOfSumKJTests and NumSquaresOfSumKDPJTests.