LC 992. K 个不同整数的子数组

题目描述

这是 LeetCode 上的 992. K 个不同整数的子数组 ，难度为 困难。

给定一个正整数数组 $A$，如果 $A$ 的某个子数组中不同整数的个数恰好为 $K$，则称 $A$ 的这个连续、不一定不同的子数组为好子数组。

例如，$[1,2,3,1,2]$ 中有 $3$ 个不同的整数：$1$，$2$，以及 $3$。

返回 $A$ 中好子数组的数目。

示例 1：

输入：A = [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]

示例 2：

输入：A = [1,2,1,3,4], K = 3

输出：3

解释：恰好由 3 个不同整数组成的子数组：[1,2,1,3], [2,1,3], [1,3,4]

提示：

• $1 <= A.length <= 20000$
• $1 <= A[i] <= A.length$
• $1 <= K <= A.length$

---

滑动窗口

对原数组每个 $nums[i]$ 而言：

1. 找到其左边「最远」满足出现 $k$ 个不同字符的下标，记为 $p$ ，这时候形成的区间为 $[p, i]$
2. 找到其左边「最远」满足出现 $k - 1$ 个不同字符的下标，记为 $j$ ，这时候形成的区间为 $[j, i]$
3. 那么对于 $j - p$ 其实就是代表以 $nums[i]$ 为右边界（必须包含 $num[i]$），不同字符数量「恰好」为 $k$ 的子数组数量

我们使用 $lower$ 数组存起每个位置的 $p$；使用 $upper$ 数组存起每个位置的 $j$。

累积每个位置的 $upper[i] - lower[i]$ 就是答案。

计算 $lower$ 数组 和 $upper$ 数组的过程可以使用双指针。

Java 代码：

class Solution {
    public int subarraysWithKDistinct(int[] nums, int k) {
        int n = nums.length, ans = 0;
        int[] lower = new int[n], upper = new int[n];
        find(nums, lower, k);
        find(nums, upper, k - 1);
        for (int i = 0; i < n; i++) ans += upper[i] - lower[i];
        return ans;
    }
    void find(int[] nums, int[] arr, int k) {
        int n = nums.length;
        int[] cnt = new int[20010];
        for (int i = 0, j = 0, sum = 0; j < n; j++) {
            if (++cnt[nums[j]] == 1) sum++;
            while (sum > k) {
                if (--cnt[nums[i++]] == 0) sum--;
            }
            if (sum == k) arr[j] = i;
        }
    }
}

C++ 代码：

class Solution {
public:
    int subarraysWithKDistinct(vector<int>& nums, int k) {
        int n = nums.size(), ans = 0;
        vector<int> lower(n, 0), upper(n, 0);
        find(nums, lower, k);
        find(nums, upper, k - 1);
        for (int i = 0; i < n; i++) ans += upper[i] - lower[i];
        return ans;
    }
    void find(vector<int>& nums, vector<int>& arr, int k) {
        int n = nums.size();
        unordered_map<int, int> cnt;
        int i = 0, j = 0, sum = 0;
        for (j = 0; j < n; j++) {
            if (++cnt[nums[j]] == 1) sum++;
            while (sum > k) {
                if (--cnt[nums[i++]] == 0) sum--;
            }
            if (sum == k) arr[j] = i;
        }
    }
};

Python 代码：

class Solution:
    def subarraysWithKDistinct(self, nums: List[int], k: int) -> int:
        n, ans = len(nums), 0
        lower, upper = [0] * n, [0] * n
        self.find(nums, lower, k)
        self.find(nums, upper, k - 1)
        return sum(upper[i] - lower[i] for i in range(n))

    def find(self, nums, arr, k):
        n = len(nums)
        cnt = defaultdict(int)
        i, j, sumv = 0, 0, 0
        while j < n:
            if cnt[nums[j]] == 0:
                sumv += 1
            cnt[nums[j]] += 1
            while sumv > k:
                cnt[nums[i]] -= 1
                if cnt[nums[i]] == 0:
                    sumv -= 1
                i += 1
            if sumv == k:
                arr[j] = i
            j += 1

• 时间复杂度：对数组进行常数次扫描。复杂度为 $O(n)$。
• 空间复杂度：$O(n)$

---

其他

这里的 $lower$ 和 $upper$ 其实可以优化掉，但也只是常数级别的优化，空间复杂度仍为 $O(n)$。

推荐大家打印一下 $lower$ 和 $upper$ 来看看，加深对 「 $upper[i] - lower[i]$ 代表了考虑 $nums[i]$ 为右边界，不同字符数量「恰好」为 $k$ 的子数组数量 」 这句话的理解。

---

最后

这是我们「刷穿 LeetCode」系列文章的第 No.992 篇，系列开始于 2021/01/01，截止于起始日 LeetCode 上共有 1916 道题目，部分是有锁题，我们将先把所有不带锁的题目刷完。

在这个系列文章里面，除了讲解解题思路以外，还会尽可能给出最为简洁的代码。如果涉及通解还会相应的代码模板。

为了方便各位同学能够电脑上进行调试和提交代码，我建立了相关的仓库：https://github.com/SharingSource/LogicStack-LeetCode 。

在仓库地址里，你可以看到系列文章的题解链接、系列文章的相应代码、LeetCode 原题链接和其他优选题解。