202503141758 930 和相同的二元子数组

和相同的二元子数组 给你一个二元数组 nums ，和一个整数 goal ，请你统计并返回有多少个和为 goal 的 非空 子数组。

子数组 是数组的一段连续部分。

示例 1：

输入： nums = [1,0,1,0,1], goal = 2 输出： 4 解释： 有 4 个满足题目要求的子数组：[1,0,1]、[1,0,1,0]、[0,1,0,1]、[1,0,1]

示例 2：

输入： nums = [0,0,0,0,0], goal = 0 输出： 15

提示：

• 1 <= nums.length <= 3 * 10^4
• nums[i] 不是 0 就是 1
• 0 <= goal <= nums.length

题解

求子数组的和为 goal 的数量。

先从DP角度考虑：

• 定义 dp[i] 为以 i 结尾的符合条件连续子数组个数
• 由于要求和为 goal，我们需要额外的状态来表示当前连续子数组和
• 即 $dp(i,j)$， 表示以 i 结尾且总和为 j 的连续子数组个数 因此结果为 $\sum dp[i][goal]$ 。

由于题目条件为 1 <= nums.length <= 3 * 10^4，故 $j<=3\ast 10^4$，整体的时间复杂度为 $O(3\ast 10^8)$ ：

int ans = 0;
for(int i = 0; i < nums.size(); i++) {
	for(int j = 0; j <= 3e4; j++) {
		dp[i][nums[i]] += 1;
		if(j >= nums[i] && i)
			dp[i][j] += dp[i-1][j-nums[i]];
	}
	ans += dp[i][goal];
}

但会超内存限制，故需优化。

首先是状态定义，有很大一部分 j 的取值范围是碰不到的，因此可以离散化存储即使用哈希表，定义 $count[j]$ 为和出现值为 j 的次数，而状态定义变为：

vector<unordered_map<int,int>> dp(nums.size());
int ans = nums[0] == goal ? 1 : 0;
dp[0][nums[0]]++;
for(int i = 1; i < nums.size(); i++) {
	dp[i][nums[i]] += 1;
	unordered_map<int,int>::iterator it;
	for(it = dp[i-1].begin(); it != dp[i-1].end(); it++)
		dp[i][nums[i] + (it->first)] += it->second;
 
	ans += dp[i][goal];
}

这样优化了一部分空间复杂度，但遇到极端数据时依然会超时。实际上 count[j] 存储的值都是 nums 数组的前缀和中的值，如下：

对于 [1,0,1,0,1]：
 
i: 0 (1, 1)
i: 1 (2, 0) (1, 1) (0, 1) // 这里的 (2，0) 是被 nums[i] + (it->first) 创建
i: 2 (2, 1) (3, 0) (1, 2)
i: 3 (1, 2) (3, 0) (2, 1) (0, 1)
i: 4 (3, 1) (4, 0) (2, 2) (1, 2)

故实际上可以使用滚动数组优化，只保留 i-1 的前缀和情况，更新到 i 只需要加上当前的 nums[i] 值并加入哈希表，即在当前 i 循环中，count 被更新前代表着以 i-1 结尾的前缀和和与其次数。

此时 count 哈希表就代替了 dp 数组的职位，在被更新加入当前的前缀和之前是 i-1 状态，在加入后就是 i 状态用于下一步的计算：

unordered_map<int, int> hash;
hash[0]++; // 由于前缀和的计算性质，需要初始化一个 {0,1} 用于统计向左取到头的情况
 
int ans = 0, current_sum = 0;
for(int i = 0; i < nums.size(); i++) {
	current_sum += nums[i];
	ans += hash[current_sum - goal];
	hash[current_sum] ++;
}

由于未更新前的 hash[goal] 等同于最开始代码中的 dp[i][goal]，故直接加到 ans 结果即可。

代码

class Solution {
public:
    int numSubarraysWithSum(vector<int>& nums, int goal) {
        unordered_map<int, int> hash;
        hash[0]++;
 
        int ans = 0, current_sum = 0;
        for(int i = 0; i < nums.size(); i++) {
            current_sum += nums[i];
            ans += hash[current_sum - goal];
            hash[current_sum] ++;
        }
        return ans;
    }
};