Leetcode 1248. 统计「优美子数组」 题解

题目链接：https://leetcode-cn.com/problems/count-number-of-nice-subarrays/


前缀和记录奇数的数量。

第一次想用（i, j）之间的循环遍历 + 剪枝：

class Solution {
public:
    int numberOfSubarrays(vector<int>& a, int k) {
        int len = a.size();
        vector<int> sum(len + 1, 0);
        for(int i = 1; i <= len; i++) {
            sum[i] = sum[i - 1] + (a[i - 1] % 2 != 0);
        }
        if(sum[len] < k) {
            return 0;
        }
        int res = 0;
        
        int i = 0, j = k;
        while(j <= len) {
            if(sum[j] - sum[i] == k) {
                res++;
            } 
            else if(sum[j] - sum[i] < k) {
                j++;
                continue;
            }
            int pos = i + 1;
            while(pos <= j - k) {
                if(sum[j] - sum[pos] == k) {
                    res++;
                } else if(sum[j] - sum[pos] < k) {
                    break;
                }
                pos++;
            }
            j++;
        }

        return res;
    }
};

但是还是超时，看到题解学到了，对于右指针 i 的前缀和 sum[i] ，只需要找到所有前缀和为 sum[i] - k 的值，所以遍历过程中每次更新前缀和及前缀和相同的次数 cnt 即可。

代码如下：

class Solution {
    vector<int> cnt;
public:
    int numberOfSubarrays(vector<int>& nums, int k) {
        int n = (int)nums.size();
        cnt.resize(n + 1, 0);
        int odd = 0, ans = 0;
        cnt[0] = 1;
        for (int i = 0; i < n; ++i) {
            odd += nums[i] & 1;
            ans += (odd >= k) ? cnt[odd - k] : 0;
            cnt[odd]++;
        }
        return ans;
    }
};