Problem description:

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

Example 1:

1
2
3
4
5
6
7
8
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.

Solution:

In this question, we can apply the idea of two sum. The difference between this question and the original two sum is, this question is an array of 1 and 0. What we can do is to find all the prefix sum and the total sum on each index i, and then find the complement in the map.

The reason to use prefix sum is because, if we want to find a subarray sum[i:j], we can solve it like this:

1
2
3
4
sum[i:j]= A[i]+ A[i+1]+ ... A[j]
        = sum[0:j]-sum[0:i-1]
        = (A[0]+ ...+ A[i]+ ... + A[j]) - (A[0]+ ...+ A[i-1])
        = sum[i:j]

One thing to notice is that, we need to put every possible complement subarray into consideration. That is the reason why we need to check the following code.

1
2
3
if(map.find(cur- S) != map.end()){
  res+= map[cur-S];
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
class Solution {
public:
    int numSubarraysWithSum(vector<int>& A, int S) {
        // find all subarray prefix sum using the idea of two sum
        // in two sum, we use unordered_map to recored every appeared number in the array
        int res= 0, cur= 0;
        unordered_map<int, int> map;
        for(int i= 0; i< A.size(); i++){
            cur+= A[i];
            if(map.find(cur- S) != map.end()){
                res+= map[cur-S];
            }
            if(cur == S) 
                res++;
            map[cur]++;
        }
        return res;
    }
};

time complexity: $O(n)$
space complexity: $O(n)$
reference:
https://goo.gl/QXnJ5U
https://goo.gl/tNeBxR