题目描述

在由若干 0 和 1 组成的数组 A 中，有多少个和为 S 的非空子数组。

样例

输入：A = [1,0,1,0,1], S = 2
输出：4
解释：
如下面黑体所示，有 4 个满足题目要求的子数组：

[1,0,1,0,1]
[1,0,1,0,1]
[1,0,1,0,1]
[1,0,1,0,1]

提示

1. A.length <= 30000
2. 0 <= S <= A.length
3. A[i] 为 0 或 1

算法1

(前缀和+哈希) $O(n)$

我们对数组求前缀和就可以在 $O(1)$ 的时间里求出区间[m, n]的和。如果我们直接枚举区间的起点和终点来求和为S的区间，那么需要 $O(n^2)$ 的时间，这可以利用哈希表来优化。我们可以只枚举终点，然后利用哈希表存储终点前出现过的前缀和及其次数，由此来得到前缀和为sum - S的前缀和的出现次数，并将当前的前缀和与该前缀和出现的次数存入哈希表中。

时间复杂度

哈希表的查询和插入都为 $O(1)$，总时间复杂度为 $O(n)$。

C++ 代码

class Solution {
public:
    int numSubarraysWithSum(vector<int>& A, int S) {
        int res = 0;
        unordered_map<int, int> hash;
        hash[0] = 1;

        for (int i = 0, sum = 0; i < A.size(); i ++ ) {
            sum += A[i];
            if (hash.count(sum - S)) res += hash[sum - S];
            hash[sum] ++ ;
        }

        return res;
    }
};

算法2

(双指针) $O(n)$

对于算法1其实我们相当于对每个终点，找到一个起点的范围[a, b]，使得该范围内的起点和当前终点构成的区间满足和为S。那么由于数组中不存在负数，随着终点向右移动，当前终点的前缀和变大，这个起点的范围也会向右移动来让起点的前缀和也变大，从而保持区间的和为S。所以我们可以利用双指针来维护当前和为S的区间。

并且由于数组中不包含负数，只有0和1，所以若区间[a, b]和当前终点构成的区间和为S，说明区间[a, b]是一段连续的0，当然如果a == b的话也可以不是0。那么双指针的做法就比较明显了，每次我们将右端点加入当前记录的区间和，接着维护区间的合法性，即不断的移动左端点直到和小于等于S。接下来我们判断是否有连续的0存在，并将答案记录下来。注意这里统计0的时候不能移动左端点，因为当右端点下一个位置也是0的时候，这一段连续的0也可以为下一个右端点服务，所以我们需要另外开一个变量统计0的个数。

C++ 代码

class Solution {
public:
    int numSubarraysWithSum(vector<int>& A, int S) {
        int res = 0;
        for (int i = 0, j = 0, sum = 0; i < A.size(); i ++ ) {
            sum += A[i];
            while (j < i && sum > S) sum -= A[j ++ ];
            if (sum == S) res ++ ;

            int k = j;
            while (k < i && A[k] == 0 && sum == S)
                k ++ , res ++ ;
        }
        return res;
    }
};