Program to count number of nice subarrays in Python

Suppose we have an array called nums and another value k. We have to find the number of nice subarrays. A subarray is said to be nice if it contains exactly k odd numbers.

So, if the input is like nums = [1,1,2,1,1], k = 3, then the output will be 2 because there are two subarrays [1,1,2,1] and [1,2,1,1] that contain exactly 3 odd numbers each.

Algorithm

To solve this problem, we will follow these steps −

• Create a list odd_indices to store indices of all odd numbers

• For each index i from 0 to length of nums − 1:

  • If nums[i] is odd, add i to odd_indices

• Use a sliding window approach with start = 0, end = k − 1

• For each valid window of k odd numbers, count all possible subarrays

• Calculate subarrays by multiplying left choices with right choices

Example

Let us see the following implementation to get better understanding −

def count_nice_subarrays(nums, k):
    # Store indices of odd numbers
    odd_indices = []
    for i in range(len(nums)):
        if nums[i] % 2 == 1:
            odd_indices.append(i)
    
    # If we don't have k odd numbers, return 0
    if len(odd_indices) < k:
        return 0
    
    start = 0
    end = k - 1
    left_boundary = 0
    count = 0
    
    # Slide the window of k odd numbers
    while end < len(odd_indices):
        # Calculate right boundary
        if end == len(odd_indices) - 1:
            right_boundary = len(nums) - 1
        else:
            right_boundary = odd_indices[end + 1] - 1
        
        # Count subarrays for current window
        left_choices = odd_indices[start] - left_boundary + 1
        right_choices = right_boundary - odd_indices[end] + 1
        count += left_choices * right_choices
        
        # Move to next window
        left_boundary = odd_indices[start] + 1
        start += 1
        end += 1
    
    return count

# Test the function
nums = [1, 1, 2, 1, 1]
k = 3
result = count_nice_subarrays(nums, k)
print(f"Number of nice subarrays: {result}")

The output of the above code is −

Number of nice subarrays: 2

How It Works

The algorithm works by:

1. Finding odd positions: First, we collect all indices where odd numbers appear in the array

2. Sliding window: We use a window of size k to consider k consecutive odd numbers

3. Counting combinations: For each window, we count how many ways we can extend left and right to form valid subarrays

Another Example

Let's test with a different input −

def count_nice_subarrays(nums, k):
    odd_indices = []
    for i in range(len(nums)):
        if nums[i] % 2 == 1:
            odd_indices.append(i)
    
    if len(odd_indices) < k:
        return 0
    
    start = 0
    end = k - 1
    left_boundary = 0
    count = 0
    
    while end < len(odd_indices):
        if end == len(odd_indices) - 1:
            right_boundary = len(nums) - 1
        else:
            right_boundary = odd_indices[end + 1] - 1
        
        left_choices = odd_indices[start] - left_boundary + 1
        right_choices = right_boundary - odd_indices[end] + 1
        count += left_choices * right_choices
        
        left_boundary = odd_indices[start] + 1
        start += 1
        end += 1
    
    return count

# Test with different inputs
test_cases = [
    ([2, 4, 6], 1),      # No odd numbers
    ([1, 3, 5], 2),      # Multiple combinations
    ([2, 2, 2, 1, 2, 2, 1, 2, 2, 2], 2)  # Sparse odd numbers
]

for nums, k in test_cases:
    result = count_nice_subarrays(nums, k)
    print(f"nums = {nums}, k = {k} ? {result}")

The output of the above code is −

nums = [2, 4, 6], k = 1 ? 0
nums = [1, 3, 5], k = 2 ? 2
nums = [2, 2, 2, 1, 2, 2, 1, 2, 2, 2], k = 2 ? 16

Conclusion

The sliding window approach efficiently counts nice subarrays by first identifying odd number positions, then using combinatorics to count all possible subarrays containing exactly k odd numbers. The time complexity is O(n) where n is the length of the input array.