题目描述

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

样例

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

限制

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

算法1

(线性扫) $O(n)$

1. 令 $f(x)$ 表示前缀和为 $x$ 的位置有多少个。注意初始时，$f(0) = 1$。
2. 我们在遍历数组的过程中，如果发现位置 $i$ 的前缀和 $c \ge S$，则当以 $i$ 为区间右端点的答案为 $f(c - S)$。
3. 每个位置需要 $f(c)$ 累加 1。

时间复杂度

• 每个位置仅遍历一次，故时间复杂度为 $O(n)$。

空间复杂度

• 需要额外 $O(n)$ 的空间存储 $f$ 数组。

C++ 代码

class Solution {
public:
    int numSubarraysWithSum(vector<int>& A, int S) {
        int n = A.size(), ans = 0;
        vector<int> f(n + 1, 0);

        f[0] = 1;
        int c = 0;
        for (int i = 0; i < n; i++) {
            c += A[i];
            if (c >= S) ans += f[c - S];
            f[c]++;
        }

        return ans;
    }
};

算法2

(双指针) $O(n)$

1. 此题的答案可以分为两部分，我们可以分别求出总和小于等于 $S$ 的区间个数，然后减去总和小于等于 $S - 1$ 的区间个数。
2. 求总和小于等于 $S$ 的区间个数可以采用双指针。
3. 对于每个位置 $i$，求出尽可能小的 $j$，满足区间 $[j, i]$ 的总和等于 $S$，此时以 $i$ 为结尾右端点的合法区间个数就是 $i - j + 1$。注意到 $j$ 是随着 $i$ 增加单调不减的。
4. 如果 $S == 0$，则在求第二部分时直接返回 $0$。

时间复杂度

• 双指针求每一部分的时间复杂度都是 $O(n)$，故总时间复杂度为 $O(n)$。

空间复杂度

• 仅需要常数的额外空间。

C++ 代码

class Solution {
public:
    int solve(const vector<int> &A, int S) {
        if (S < 0) return 0;
        int n = A.size(), res = 0;

        for (int i = 0, j = 0, c = 0; i < n; i++) {
            c += A[i];
            while (c > S) {
                c -= A[j];
                j++;
            }

            res += i - j + 1;
        }

        return res;
    }

    int numSubarraysWithSum(vector<int>& A, int S) {
        return solve(A, S) - solve(A, S - 1);
    }
};