1248. Count Number of Nice Subarrays

Given an array of integers nums and an integer k. A continuous subarray is called nice if there are k odd numbers on it.

Return the number of nice sub-arrays.

 

Example 1:

Input: nums = [1,1,2,1,1], k = 3
Output: 2
Explanation: The only sub-arrays with 3 odd numbers are [1,1,2,1] and [1,2,1,1].

Example 2:

Input: nums = [2,4,6], k = 1
Output: 0
Explanation: There is no odd numbers in the array.

Example 3:

Input: nums = [2,2,2,1,2,2,1,2,2,2], k = 2
Output: 16

 

Constraints:

  • 1 <= nums.length <= 50000
  • 1 <= nums[i] <= 10^5
  • 1 <= k <= nums.length

Rust Solution

struct Solution;

impl Solution {
    fn number_of_subarrays(nums: Vec<i32>, k: i32) -> i32 {
        let n = nums.len();
        let mut count = vec![0; n + 1];
        let k = k as usize;
        count[0] = 1;
        let mut prev = 0;
        let mut res = 0;
        for i in 0..n {
            if nums[i] % 2 == 1 {
                prev += 1;
            }
            count[prev] += 1;
            if prev >= k {
                res += count[prev - k];
            }
        }
        res as i32
    }
}

#[test]
fn test() {
    let nums = vec![1, 1, 2, 1, 1];
    let k = 3;
    let res = 2;
    assert_eq!(Solution::number_of_subarrays(nums, k), res);
    let nums = vec![2, 4, 6];
    let k = 1;
    let res = 0;
    assert_eq!(Solution::number_of_subarrays(nums, k), res);
    let nums = vec![2, 2, 2, 1, 2, 2, 1, 2, 2, 2];
    let k = 2;
    let res = 16;
    assert_eq!(Solution::number_of_subarrays(nums, k), res);
}

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