Count Number of Nice Subarrays [Medium] — atMost(k)−atMost(k−1)

Sanjeev SharmaSanjeev Sharma
1 min read

Advertisement

Problem Statement

Return the number of subarrays with exactly k odd numbers.

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

Intuition

exactly(k) = atMost(k) - atMost(k-1).

atMost(k): sliding window counting odd numbers; shrink when count > k; add right-left+1 for each right.


Solutions

C++

int atMost(vector<int>& nums, int k) {
    int left=0, odds=0, res=0;
    for(int r=0;r<nums.size();r++){
        if(nums[r]%2) odds++;
        while(odds>k){if(nums[left++]%2)odds--;}
        res+=r-left+1;
    }
    return res;
}
int numberOfSubarrays(vector<int>& nums, int k){return atMost(nums,k)-atMost(nums,k-1);}

Python

def numberOfSubarrays(nums: list[int], k: int) -> int:
    def at_most(k):
        left = odds = res = 0
        for right, v in enumerate(nums):
            if v % 2: odds += 1
            while odds > k:
                if nums[left] % 2: odds -= 1
                left += 1
            res += right - left + 1
        return res
    return at_most(k) - at_most(k-1)

Complexity

  • Time: O(n), Space: O(1)

Advertisement

Sanjeev Sharma

Written by

Sanjeev Sharma

Full Stack Engineer · E-mopro