Taro Logo

Count Number of Nice Subarrays

Medium
a month ago

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.

For example:

  • nums = [1,1,2,1,1], k = 3 should return 2 because the only sub-arrays with 3 odd numbers are [1,1,2,1] and [1,2,1,1].
  • nums = [2,4,6], k = 1 should return 0 because there are no odd numbers in the array.
  • nums = [2,2,2,1,2,2,1,2,2,2], k = 2 should return 16.

How would you approach this problem?

Sample Answer
class Solution:
    def numberOfSubarrays(self, nums: list[int], k: int) -> int:
        def atMost(nums: list[int], k: int) -> int:
            if k < 0:
                return 0
            left = 0
            count = 0
            res = 0
            for right in range(len(nums)):
                if nums[right] % 2 != 0:
                    count += 1
                while count > k:
                    if nums[left] % 2 != 0:
                        count -= 1
                    left += 1
                res += right - left + 1
            return res

        return atMost(nums, k) - atMost(nums, k - 1)

Explanation

The problem requires finding the number of continuous subarrays that contain exactly k odd numbers. A sliding window approach can efficiently solve this.

  1. Naive Approach: Iterate through all possible subarrays and count odd numbers in each. This approach has a time complexity of O(n^2), where n is the length of the input array.

  2. Optimal Approach: Use a sliding window technique. The function atMost(nums, k) computes the number of subarrays with at most k odd numbers. Then, subtract the number of subarrays with at most k-1 odd numbers to get the exact number of subarrays with k odd numbers. This reduces time complexity to O(n).

  • The outer loop iterates through the array using the right pointer.
  • When an odd number is encountered, the count of odd numbers in the current window increases.
  • If count exceeds k, the left pointer moves to shrink the window until the count is reduced to at most k.
  • The number of valid subarrays ending at the right pointer is right - left + 1, which is added to res.

Example

Consider nums = [1,1,2,1,1], k = 3

  1. Call atMost(nums, 3): this will count all subarrays with at most 3 odd numbers.
  2. Call atMost(nums, 2): this will count all subarrays with at most 2 odd numbers.
  3. Subtract the result of step 2 from step 1. The result is the number of subarrays with exactly 3 odd numbers.

Big(O) Run-time Analysis

The numberOfSubarrays function calls atMost twice, and atMost iterates through the array once. Thus, the time complexity is O(n), where n is the length of the input array.

Big(O) Space Usage Analysis

The space complexity is O(1) as no extra space that scales with the input is used. Only constant extra space is used for variables like left, count, and res.

Edge Cases

  1. Empty Array: If the input array is empty, the function will correctly return 0 because there are no subarrays.
  2. k = 0: If k is 0, the function calculates subarrays with no odd numbers.
  3. k > number of odd numbers in nums: The function correctly handles cases where k is greater than the number of odd numbers in the input array.
  4. All even numbers: If the input array contains only even numbers and k > 0, the function will return 0 since no subarray can contain k odd numbers.
  5. All odd numbers: If the input array contains only odd numbers, the function will return the correct number of subarrays with exactly k odd numbers.