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.

Rust Solution

struct Solution;

impl Solution {
    fn num_subarrays_with_sum(a: Vec<i32>, s: i32) -> i32 {
        let n = a.len();
        let mut count = vec![0; n + 1];
        count[0] = 1;
        let mut sum = 0;
        let mut res = 0;
        for i in 0..n {
            sum += a[i];
            if sum >= s {
                res += count[(sum - s) as usize];
            }
            count[sum as usize] += 1;
        }
        res as i32
    }
}

#[test]
fn test() {
    let a = vec![1, 0, 1, 0, 1];
    let s = 2;
    let res = 4;
    assert_eq!(Solution::num_subarrays_with_sum(a, s), res);
}

Having problems with this solution? Click here to submit an issue on github.