leetcode 周赛 161 5248. 统计「优美子数组」 1248. Count Number of Nice Subarrays

目录

1. 解答

https://leetcode-cn.com/contest/weekly-contest-161/problems/count-number-of-nice-subarrays/

给你一个整数数组 nums 和一个整数 k。

如果某个子数组中恰好有 k 个奇数数字，我们就认为这个子数组是「优美子数组」。

请返回这个数组中「优美子数组」的数目。

示例 1：

输入：nums = [1,1,2,1,1], k = 3
输出：2
解释：包含 3 个奇数的子数组是 [1,1,2,1] 和 [1,2,1,1] 。

示例 2：

输入：nums = [2,4,6], k = 1
输出：0
解释：数列中不包含任何奇数，所以不存在优美子数组。

示例 3：

输入：nums = [2,2,2,1,2,2,1,2,2,2], k = 2
输出：16

提示：

• 1 <= nums.length <= 50000
• 1 <= nums[i] <= 10^5
• 1 <= k <= nums.length

解答

我的方法还是有点麻烦的，但是思路清晰，nx 统计了下一个奇数的位置，pr 统计了上一个奇数的位置。

然后我们设定一个闭区间 [st, ed]。使 st 和 ed 对应的数都是奇数。那么这个区间对应的答案数是，st到前一个奇数或数组开始的距离 乘以 ed 到下一个奇数或者结尾的距离。

class Solution:
    def numberOfSubarrays(self, nums: List[int], k: int) -> int:
        nx = [-1 for i in range(len(nums))]
        pr = [-1 for i in range(len(nums))]
        last = -1
        for i in range(len(nums)):
            pr[i] = last
            if nums[i] % 2 == 1:
                last = i
        last = -1
        for i in range(len(nums)-1, -1, -1):
            nx[i] = last
            if nums[i] % 2 == 1:
                last = i
        cnt = 1
        if nums[0] % 2 == 1:
            st = 0
            ed = 0
        else:
            st = nx[0]
            ed = nx[0]
        ans = 0
        while st != -1 and ed != -1:
            if cnt == k:
                if pr[st] == -1: p_l = st + 1
                else: p_l = st - pr[st]
                if nx[ed] == -1: n_l = len(pr) - ed
                else: n_l = nx[ed] - ed
                #print(st, ed, p_l, n_l)
                ans += p_l * n_l
            if cnt < k:
                ed = nx[ed]
                cnt += 1
            else:
                st = nx[st]
                cnt -= 1
        #print(pr)
        #print(nx)
        return ans