统计所有连续子数组，使其元素之和为零

Question

给定一个长度为n的随机数数组(包括正数和负数)，我想要找到相加等于零的子数组个数。

示例: 给定数组a={1, -1 ,2 , -2, 6, -6}，输出结果应为6，因为子数组如下：

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

我知道暴力解决方案的时间复杂度是O(n^2)，我想要另一种时间复杂度为O(n)或者O(n log n)的解决方案，你能帮忙提供吗？

- Mohamed Yakout

---

我不认为“连续子数组求和”问题有“O(n)或O(n log n)”的解决方案。 - Haris

由于可能的子数组总数为O(n^2)，我怀疑你不会想出比这更好的通用算法。如果你确实有，那么它必须是一种不仅仅枚举子数组的方法。考虑这个困难的情况：{6,0,0,...0,0,-6}。 - RBarryYoung

考虑 { 0, 0, ... 0, 0 }，现在所有的东西都加起来等于零，你必须输出平方数量级的东西。 - harold

FYI：我认为下面的@DavidEisenstat证明了我的错误。 - RBarryYoung

Answer 1

设数组为a1，...，an。令s0，...，sn是由sj = a1 + ... + aj（其中s0 = 0）定义的前缀和。从i到j（包括i和j）的子数组之和是sj - si-1。使用映射实现O(n)/O(n log n)时间复杂度算法，计算每个前缀和中每个数字出现的次数。对于该映射中的值k，计算k个选择2的和。

例如，如果我们有以下数组

1 -1 2 -2 6 -6

那么前缀和为：

0 1 0 2 0 6 0

并且计数如下

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

输出结果为4选2 + 1选2 + 1选2 + 1选2 = 6（其中k选2 = k(k-1)/2）。

在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())

- David Eisenstat

---

+1：哇，不错！我没想到有绕过枚举它们的方法，但这个方法似乎很有效。 - RBarryYoung

我不明白在获取前缀和之后下一步是什么？ - Mohamed Yakout

@DavidEisenstat 我不明白 sum(v*(v-1) // 2 for v in freq.values()) 的原因。 （翻译：I don't understand the reason behind sum(v*(v-1) // 2 for v in freq.values()) that you mentioned to DavidEisenstat.） - Shubham Aggarwal

一切都没问题。但是最后一行“输出为4选2 + 1选2 + 1选2 + 1选2 = 6（其中k选2 = k(k-1)/2）”的意思是什么？这不清楚。 - Imtiaz Mehedi