K个不同子数组的数目--滑动窗口--字节--亚马逊

Stay hungry, stay foolish

题目描述

给定一个正整数数组 nums和一个整数 k，返回 nums 中 「好子数组」 的数目。 如果 nums 的某个子数组中不同整数的个数恰好为 k，则称 nums 的这个连续、不一定不同的子数组为 「好子数组 」。 例如，[1,2,3,1,2] 中有 3 个不同的整数：1，2，以及 3。 子数组 是数组的 连续 部分。

输入：nums = [1,2,1,2,3], k = 2
输出：7
解释：恰好由 2 个不同整数组成的子数组：[1,2], [2,1], [1,2], [2,3], [1,2,1], [2,1,2], [1,2,1,2].

输入：nums = [1,2,1,3,4], k = 3
输出：3
解释：恰好由 3 个不同整数组成的子数组：[1,2,1,3], [2,1,3], [1,3,4].

思路分析

问题分解：

• 我们想要找到所有包含恰好 k 个不同元素的子数组。

• 但是直接计算恰好 k 个不同元素的子数组比较麻烦，所以我们可以换个思路：先计算所有包含 至多 k 个不同元素的子数组数量，再减去所有包含 至多 k-1 个不同元素的子数组数量。这样就得到了恰好 k 个不同元素的子数组数量。

函数 f(nums, k)：

• 这个函数的作用是计算所有包含 至多 k 个不同元素的子数组数量。

• 使用两个指针 l 和 r，分别表示当前窗口的左边界和右边界。

• 使用一个数组 cnts 来记录当前窗口中每个数字的出现次数。

• 使用一个变量 collect 来记录当前窗口中有多少个不同的数字。

滑动窗口的工作原理：

• 右指针 r 不断向右移动，每遇到一个新的数字，就更新 cnts 和 collect。

• 如果 collect 超过了 k，说明当前窗口中的不同数字太多了，需要收缩左指针 l，直到 collect 不再超过 k。

• 每次右指针 r 移动后，当前窗口中所有以 l 为起点、r 为终点的子数组都是有效的（即包含至多 k 个不同元素），所以把这些子数组的数量累加到结果 ans 中。

最终结果：

• 通过调用 f(nums, k) 和 f(nums, k - 1)，我们得到了包含 至多 k 个不同元素的子数组数量和包含 至多 k-1 个不同元素的子数组数量。

• 两者相减，就得到了包含 恰好 k 个不同元素的子数组数量。

1. subarraysWithKDistinct 方法：

• return f(nums, k) - f(nums, k - 1);：计算包含恰好 k 个不同元素的子数组数量。通过计算包含至多 k 个不同元素的子数组数量，减去包含至多 k-1 个不同元素的子数组数量，可以得到恰好 k 个不同元素的子数组数量。

2. f 方法：

• int ans = 0;：初始化结果变量 ans 为 0，用于存储当前窗口中所有子数组的数量。

• vector cnts(20001, 0);：初始化一个大小为 20001 的计数数组 cnts，用于记录每个元素的出现次数。假设输入数组中的元素范围在 0 到 20000 之间。

• for (int l = 0, r = 0, collect = 0; r < arr.size(); r++) {：初始化左指针 l 和右指针 r，以及一个变量 collect 用于记录当前窗口中不同元素的数量。右指针 r 从 0 开始遍历数组。

• if (++cnts[arr[r]] == 1) { collect++; }：右指针 r 向右移动，并更新计数数组 cnts。如果当前元素 arr[r] 的计数从 0 变为 1，说明这是一个新出现的不同元素，collect 增加 1。

• while (collect > k) { if (--cnts[arr[l++]] == 0) { collect--; } }：如果当前窗口中不同元素的数量超过 k，左指针 l 向右移动，直到窗口中不同元素的数量不超过 k。在移动左指针 l 时，更新计数数组 cnts，如果某个元素的计数从 1 变为 0，说明这个元素不再出现在当前窗口中，collect 减少 1。

• ans += r - l + 1;：当前窗口中所有子数组的数量 (从 l 到 r) 都是有效的，累加到结果 ans 中。

• return ans;：返回结果 ans。

完整代码

class Solution {
public:
    int subarraysWithKDistinct(vector<int>& nums, int k) {
        // 计算包含恰好 k 个不同元素的子数组数量
        // 通过计算包含至多 k 个不同元素的子数组数量，减去包含至多 k-1 个不同元素的子数组数量
        return f(nums, k) - f(nums, k - 1);
    }

    int f(vector<int> arr, int k) {
        // 初始化结果变量 ans 为 0
        int ans = 0;
        // 初始化一个大小为 20001 的计数数组 cnts，用于记录每个元素的出现次数
        vector<int> cnts(20001, 0);
        // 初始化左指针 l 和右指针 r，以及一个变量 collect 用于记录当前窗口中不同元素的数量
        for (int l = 0, r = 0, collect = 0; r < arr.size(); r++) {
            // 右指针 r 向右移动，并更新计数数组和不同元素的数量
            if (++cnts[arr[r]] == 1) {
                collect++;
            }
            // 如果当前窗口中不同元素的数量超过 k，左指针 l 向右移动，直到窗口中不同元素的数量不超过 k
            while (collect > k) {
                if (--cnts[arr[l++]] == 0) {
                    collect--;
                }
            }
            // 当前窗口中所有子数组的数量 (从 l 到 r) 都是有效的，累加到结果 ans 中
            ans += r - l + 1;
        }
        // 返回结果
        return ans;
    }
};