题目描述

给定一个只包含0和1的数组A，求它共有多少个子区间的和是S？

注意：

1. A.length <= 30000
2. 0 <= S <= A.length
3. A[i]不是0就是1

样例

输入：A = [1,0,1,0,1], S = 2
输出：4
解释：
四个子区间分别是：
[0, 2]
[0, 3]
[1, 4]
[2, 4]

算法

(前缀和) $O(n)$

最暴力的做法是依次枚举区间的起点、终点，然后求区间总和，时间复杂度是 $O(n^3)$。

然后我们考虑一下怎么优化？

1. 可以预处理出数组前缀和，这样求区间和时可以用 $O(1)$ 的时间算出，比如区间[i, j]的和是sum[j] - sum[i - 1]。时间复杂度降为 $O(n^2)$。
2. 我们可以从前往后枚举区间终点，同时用一个数组记录当前不同前缀和的数量，比如f[i]表示和为i的前缀个数。假设枚举的当前坐标是j，那么我们的目标就是计算j之前共有多少个前缀和是sum[j] - S，这个值就是f[sum[j] - S]；

最后将所有终点对应的区间个数累加起来，就是答案。

时间复杂度分析：

• 求前缀和的时间复杂度是 $O(n)$；
• 遍历区间终点，同时更新 f[i] 的时间复杂度也是 $O(n)$；

所以总时间复杂度是 $O(n)$。

C++ 代码

class Solution {
public:
    int numSubarraysWithSum(vector<int>& A, int S) {
        int n = A.size();
        vector<int> sum(n + 1, 0), f(n + 1, 0); // f[i] 记录前面有多少个和为i的前缀
        for (int i = 1; i <= n; i ++ ) sum[i] = sum[i - 1] + A[i - 1];
        f[0] = 1;
        int res = 0;
        for (int i = 1; i <= n; i ++ )
        {
            int s = sum[i];
            if (s >= S) res += f[s - S];
            f[s] ++ ;
        }
        return res;
    }
};