Taro Logo

Count Number of Nice Subarrays

Medium
Amazon logo
Amazon
0 views
Topics:
ArraysSliding WindowsTwo Pointers

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 are 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

Solution


Naive Approach

One straightforward approach is to generate all possible subarrays and check if each subarray is "nice" (i.e., contains exactly k odd numbers). This brute force method involves iterating through all possible start and end indices of subarrays and counting the odd numbers within each.

def number_of_subarrays_brute_force(nums, k):
    count = 0
    for i in range(len(nums)):
        for j in range(i, len(nums)):
            subarray = nums[i:j+1]
            odd_count = 0
            for num in subarray:
                if num % 2 != 0:
                    odd_count += 1
            if odd_count == k:
                count += 1
    return count

Time Complexity: O(n^3), where n is the length of the input array nums. This is because we have nested loops to generate all subarrays (O(n^2)), and another loop to count odd numbers in each subarray (O(n)).

Space Complexity: O(n), where n is the length of the longest subarray. In the worst case, we store a complete copy of the input array.

Optimal Approach: Sliding Window

A more efficient approach is to use a sliding window technique. The core idea is to maintain a window and adjust its boundaries to ensure that the window always contains exactly k odd numbers. We can then count the number of such windows.

Instead of directly counting subarrays with exactly k odd numbers, we can transform the problem. The number of subarrays with exactly k odd numbers can be found by subtracting the number of subarrays with at most k - 1 odd numbers from the number of subarrays with at most k odd numbers.

def at_most(nums, k):
    if k < 0:
        return 0
    
    count = 0
    left = 0
    odd_count = 0
    
    for right in range(len(nums)):
        if nums[right] % 2 != 0:
            odd_count += 1
        
        while odd_count > k:
            if nums[left] % 2 != 0:
                odd_count -= 1
            left += 1
        
        count += right - left + 1
    
    return count


def number_of_subarrays(nums, k):
    return at_most(nums, k) - at_most(nums, k - 1)

Explanation

  1. at_most(nums, k) calculates the number of subarrays with at most k odd numbers.
  2. The outer loop iterates through the array using the right pointer, expanding the window.
  3. If nums[right] is odd, increment odd_count.
  4. The while loop contracts the window from the left (left pointer) when odd_count exceeds k.
  5. count += right - left + 1 adds the number of valid subarrays ending at the current right index.

Time Complexity: O(n), where n is the length of the input array nums. We iterate through the array at most twice (once with the right pointer and once with the left pointer).

Space Complexity: O(1). We use only a constant amount of extra space for variables like count, left, right, and odd_count.

Edge Cases:

  • If k is 0 and there are no odd numbers, the correct answer is 0.
  • If the array contains no odd numbers and k > 0, the answer is 0.
  • If k is larger than the number of odd numbers in the array, then we need to handle accordingly. Our at_most function implicitly handles this because if k > total number of odd numbers in array, then left will never increment and right will traverse the entire array, counting all possible subarrays.