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:
A.length <= 30000
0 <= S <= A.length
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);
}