Counting all contiguous sub-arrays given sum zero

前端 未结 1 1441
不思量自难忘°
不思量自难忘° 2021-01-02 23:52

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

Example: Given I h

相关标签:
1条回答
  • 2021-01-03 00:16

    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())
    
    0 讨论(0)
提交回复
热议问题