leetcode - 1248. Count Number of Nice Subarrays

Description

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

Solution

Presum

Solved after hints.

Change all the odd numbers into 1s, and even numbers into 0s. Then this problem can be solved by prefix sum.

The key is the current sum, and the value is a list, where it stores [start_index, end_index]. So when cur_sum - k appears in the dict, that means we have end_index - start_index + 1 valid solutions.

Time complexity: $o(n)$
Space complexity: $o(n)$

Sliding window

Similar to 992. K个不同整数的子数组, k odd numbers is equal to at most k odd numbers - at most k-1 odd numbers.

Time complexity: $o(n)$
Space complexity: $o(1)$

Code

Presum

class Solution:
    def numberOfSubarrays(self, nums: List[int], k: int) -> int:
        for i in range(len(nums)):
            nums[i] &= 1
        pre_sum = {0: [-1, -1]}
        cur_sum = 0
        res = 0
        for i, each_num in enumerate(nums):
            cur_sum += each_num
            if cur_sum - k in pre_sum:
                res += pre_sum[cur_sum - k][1] - pre_sum[cur_sum - k][0] + 1
            if cur_sum in pre_sum:
                pre_sum[cur_sum][1] = i
            else:
                pre_sum[cur_sum] = [i, i]
        return res

Sliding window

class Solution:
    def numberOfSubarrays(self, nums: List[int], k: int) -> int:
        left1, left2 = 0, 0
        cnt1, cnt2 = 0, 0
        res = 0
        for right in range(len(nums)):
            if nums[right] & 1 == 1:
                cnt1 += 1
                cnt2 += 1
            while cnt1 > k:
                if nums[left1] & 1 == 1:
                    cnt1 -= 1
                left1 += 1
            while cnt2 > k - 1:
                if nums[left2] & 1 == 1:
                    cnt2 -= 1
                left2 += 1
            res += left2 - left1
        return res