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.

题目理解：

给定一个01数组，和一个指定数S，求出有多少个子数组的和是S

解题思路：

找到以1开始并且以1结尾，并且和为S的子串，然后计算这个子串的左边有多少个0，右边有多少个0，相乘再相加。

因为元素只有0和1，所以找子串的过程可以是线性的，每次放弃最左边的1，并且在最右边添加一个1就行

代码如下：

class Solution {
    public int numSubarraysWithSum(int[] A, int S) {
        if(S == 0){
            int count = 0, res = 0;
            for(int i : A){
                if(i == 1){
                    res += count * (count + 1) / 2;
                    count = 0;
                }
                else
                    count++;
            }
            res += count * (count + 1) / 2;
            return res;
        }
        int sum = 0;
        int left = 0, right = 0;
        int len = A.length;
        int res = 0;
        while(right < len){
            while(right < len && sum < S){
                sum += A[right++];
            }
            if(right == len && sum < S)
                break;
            int front = 1, back = 1;
            while(left < right && A[left] == 0){
                front++;
                left++;
            }
                
            while(right < len && A[right] == 0){
                back++;
                right++;
            }
            sum -= A[left++];
            res += front * back;
        }
        return res;
    }
}