Counting all contiguous subarrays specified by the sum of zero
Given an array of random numbers (positive and negative) with length n, I want the number of contiguous subarrays to be zero.
Example:
Given that I have an array a={1, -1 ,2 , -2, 6, -6}
, the output will be 6
because the subarrays are as follows:
1 -1 & 2 -2 & 6 -6 & 1 -1 2 -2 & 2 -2 6 -6 & 1 -1 2 -2 6 -6
I know a brute force O (n ^ 2) solution, I want any other O (n) or O (n log n) solution?
source to share
Let the array be a1, ..., an. Let s0, ..., sn be the prefix sums defined by sj = a1 + ... + aj (and s0 = 0). The sum of the submatrix from i to j inclusive is sj - si-1. For an O (n) / O (n log n) -time algorithm, using a map, count the number of occurrences of each number among the prefix sums. The sum of k selects 2 for k in the values of this mapping.
For example, if we have an array
1 -1 2 -2 6 -6
then the prefix sums
0 1 0 2 0 6 0
and counting
0: 4
1: 1
2: 1
3: 1
and output 4 chooses 2 + 1 chooses 2 + 1 chooses 2 + 1 chooses 2 = 6 (where k chooses 2 = k (k-1) / 2).
In Python:
def countsumzero(lst):
prefixsums = [0]
for x in lst:
prefixsums.append(prefixsums[-1] + x)
freq = {}
for y in prefixsums:
if y in freq:
freq[y] += 1
else:
freq[y] = 1
return sum(v*(v-1) // 2 for v in freq.values())
source to share