C++用哈希表统计和为0的子数组时使用nC2组合公式的原因

核心逻辑前提 #

子数组nums[i...j]（下标从i到j，闭区间）的和可以用前缀和快速计算：

子数组和 = 前缀和[j] - 前缀和[i-1]
我们要求子数组和为0，等价于 前缀和[j] == 前缀和[i-1]，其中前缀和[k]表示数组前k个元素的累加和，通常会额外定义前缀和[0] = 0（表示前0个元素和为0，用于匹配从数组开头到某个位置的和为0的情况）。

组合公式nC2的作用 #

nC2是组合数公式，含义是从n个不同元素中任选2个的方案总数，计算式为 n*(n-1)/2。
在该算法中，它的作用是快速统计同一个前缀和值对应的、和为0的子数组总数：

• 如果某个前缀和值一共出现了n次，说明有n个不同的下标对应这个前缀和
• 从这n个下标中任选2个，设为a和b（a < b），那么子数组nums[a ... b-1]的和一定为0
• 所有满足条件的子数组总数就是这n个下标选2个的组合数，即nC2

对应示例计算 #

你给出的数组是{1,-1,1,-1}，完整的前缀和序列（包含初始的前缀和[0]=0）为{0,1,0,1,0}：

• 前缀和为0的出现次数是3次，贡献的子数组数量为 3*2/2 = 3
• 前缀和为1的出现次数是2次，贡献的子数组数量为 2*1/2 =1
• 总符合条件的子数组数量为3+1=4，和实际枚举的结果一致。

C++实现示例 #

#include <iostream>
#include <vector>
#include <unordered_map>
using namespace std;

int subarraySumZero(vector<int>& nums) {
    unordered_map<int, int> preSumCnt;
    preSumCnt[0] = 1; // 初始前缀和0出现1次
    int preSum = 0, res = 0;
    for (int num : nums) {
        preSum += num;
        preSumCnt[preSum]++;
    }
    // 遍历每个前缀和的出现次数，累加组合数
    for (auto& pair : preSumCnt) {
        int cnt = pair.second;
        res += cnt * (cnt - 1) / 2;
    }
    return res;
}

int main() {
    vector<int> nums = {1,-1,1,-1};
    cout << subarraySumZero(nums) << endl; // 输出4
    return 0;
}

也可以在遍历前缀和的同时直接计算结果，不用二次遍历哈希表：

int subarraySumZero_opt(vector<int>& nums) {
    unordered_map<int, int> preSumCnt;
    preSumCnt[0] = 1;
    int preSum = 0, res = 0;
    for (int num : nums) {
        preSum += num;
        // 之前已经出现过preSum的次数就是本次新增的符合条件的子数组数
        res += preSumCnt[preSum];
        preSumCnt[preSum]++;
    }
    return res;
}

两种实现的时间复杂度都是O(n)，空间复杂度最坏为O(n)，远优于暴力枚举的O(n²)时间复杂度。

内容的提问来源于stack exchange，提问作者Anonymous User