930. Binary Subarrays With Sum

In an array A of 0s and 1s, how many non-empty subarrays have sum S?

Example 1:

Input: A = [1,0,1,0,1], S = 2
Output: 4
Explanation: 
The 4 subarrays are bolded below:
[1,0,1,0,1]
[1,0,1,0,1]
[1,0,1,0,1]
[1,0,1,0,1]

Note:

  1. A.length <= 30000

  2. 0 <= S <= A.length

  3. A[i] is either 0 or 1.

// Pre-Sum
int numSubarraysWithSum(vector<int>& A, int S) { // time: O(n); space: O(n)
    unordered_map<int, int> mp;
    ++mp[0];
    int presum = 0, res = 0;
    for (int num : A) {
        presum += num;
        res += mp[presum - S];
        ++mp[presum];
    }
    return res;
}
// Presum
int numSubarraysWithSum(vector<int>& A, int S) { // time: O(n); space: O(n)
    vector<int> cnt(A.size() + 1);
    int sum = 0, res = 0;
    for (int num : A) {
        sum += num;
        if (sum == S) ++res;
        if (sum - S >= 0) res += cnt[sum - S];
        ++cnt[sum];
    }
    return res;
}
// Sliding Window
int atMost(const vector<int>& A, int S) {
    if (S < 0) return 0;
    int res = 0, i = 0, n = A.size();
    for (int j = 0; j < n; ++j) {
        S -= A[j];
        while (S < 0) {
            S += A[i++];
        }
        res += j - i + 1;
    }
    return res;
}
int numSubarraysWithSum(vector<int>& A, int S) { // time: O(n); space: O(n)
    return atMost(A, S) - atMost(A, S - 1);
}

Last updated

Was this helpful?