Find number of continuous subarray having sum zero

Consider S[0..N] - prefix sums of your array, i.e. S[k] = A[0] + A[1] + ... + A[k-1] for k from 0 to N.

Now sum of elements from L to R-1 is zero if and only if S[R] = S[L]. It means that you have to find number of indices 0 <= L < R <= N such that S[L] = S[R].

This problem can be solved with a hash table. Iterate over elements of S[] while maintaining for each value X number of times it was met in the already processed part of S[]. These counts should be stored in a hash map, where the number X is a key, and the count H[X] is the value. When you meet a new elements S[i], add H[S[i]] to your answer (these account for substrings ending with (i-1)-st element), then increment H[S[i]] by one.

Note that if sum of absolute values of array elements is small, you can use a simple array instead of hash table. The complexity is linear on average.

Here is the code:

long long CountZeroSubstrings(vector<int> A) {
    int n = A.size();

    vector<long long> S(n+1, 0);
    for (int i = 0; i < n; i++)
        S[i+1] = S[i] + A[i];

    long long answer = 0;
    unordered_map<long long, int> H;
    for (int i = 0; i <= n; i++) {
        if (H.count(S[i]))
            answer += H[S[i]];
        H[S[i]]++;      
    }

    return answer;
}

Tags:

Algorithm