How to check sums of all possible rectangles of array

Question

Let us assume that we have a two dimensional array A (n X n). All elements of A are either O or 1. We also have a given integer K. Our task is to find the number of all possible "rectangles" in A, which contain elements with total sum K.

To give an example , if A = 
0 0 1 0 
1 0 0 1
1 1 1 1
1 0 0 1   and k=3 ,

0 0 1 0
1 0 0 1  holds the property ,

1 1 1 holds the property ,

  1 1 1 holds the property ,

0 0 
1 0 
1 1 holds the property ,

1 1
1 0  holds the property ,

    1 1
    0 1  holds the property ,

1
1
1  holds the property

1
1
1 holds the property 

So unless I missed something, the answer should be 8 for this example.

In other words, we need to check all possible rectangles in A to see if the sum of their elements is K. Is there a way to do it faster than O(n^2 * k^2) ?

Answer 1

You could do this in O(n^3).

First note that a summed area table allows you to compute the sum of any rectangle in O(1) time given O(n^2) preprocessing time.

In this problem we only need to sum the columns, but the general technique is worth knowing.

Then for each start row and end row combination you can do a linear scan across the matrix to count the solutions either with a two pointers approach or simply by storing the previous sums.

Example Python code (finds 14 solutions to your example):

from collections import defaultdict
A=[[0, 0, 1, 0],
[1, 0, 0, 1],
[1, 1, 1, 1],
[1, 0, 0, 1]]
k=3

h=len(A)
w=len(A[0])

C=[ [0]*w for i in range(h+1)]
for x in range(w):
    for y in range(1,h+1):
        C[y][x] = C[y-1][x] + A[y-1][x]
# C[y][x] contains sum of all values A[y2][x] with y2<y

count=0
for start_row in range(h):
    for end_row in range(start_row,h):
        D=defaultdict(int) # Key is sum of columns from start to here, value is count
        D[0]=1
        t=0 # Sum of all A[y][x] for x <= col, start_row<=y<=end_row
        for x in range(w):
            t+=C[end_row+1][x] - C[start_row][x]
            count += D[t-k]
            D[t] += 1
print count