Counting all contiguous sub-arrays given sum zero

Question

Given array of random numbers(positive, and negative) with length n, I want number contiguous sub-arrays which sum equals zero.

Example: Given I have array a={1, -1 ,2 , -2, 6, -6}, output will be 6 because sub-arrays are as the following:

1 -1 & 2 -2 & 6 -6 & 1 -1 2 -2 & 2 -2 6 -6 & 1 -1 2 -2 6 -6

I know brute force solution O(n^2), I want any another solution O(n), or O(n log n) ?

Answer 1

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 subarray 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. Sum k choose 2 for k in the values of this map.

For example, if we have your array

1 -1 2 -2 6 -6

then the prefix sums are

0 1 0 2 0 6 0

and the counts are

0: 4
1: 1
2: 1
3: 1

and the output is 4 choose 2 + 1 choose 2 + 1 choose 2 + 1 choose 2 = 6 (where k choose 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())