每日算法系列【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

题解

这题最暴力的方法就是用一个字典维护每个数出现的次数，然后遍历所有的区间，求出不同整数个数正好等于 K 的区间个数。但是这种方法时间复杂度是 $O(n^2)$ ，一定会超时，所以考虑其他方法。

现在考虑右边界为 j 的情况，左边界 i 有什么规律呢？我们可以证明，满足 [i, j] 正好包含 K 个不同整数的 i 的取值是一段连续的区间。假设 [i, j]包含 K 个不同整数，同时 [i', j] 也包含 K 个不同整数（i < i'），因为从 i 移动到 i' 每个数的数量是非增的，所以这过程中没有增加新的数，也没有任何一个数的数量降到了0。

有了这个性质之后，对于任意的 j ，我们只需要求出左边界 i 的取值范围就行了。同样这里还是不能暴力求，不然就和一开始没区别了嘛。既然这样，想想如果 j 的左边界 i 的范围得到了，这时候我们继续求 j + 1 的左边界范围，能不能利用一下之前得到的结果？而不用重新计算。很容易发现，如果 j 右移了， i 的取值范围也会右移，因为 j 右移有两种结果：一是引入了新的数，二是某个存在的数的数量加 1 。第一种情况对左边界没有任何影响，因为不同整数数量没有变化，还是 K 。第二种情况不同整数数量变成 K + 1 了，这时候左边界一定要右移，删掉点数，才可能使区间符合题意。

有了上述的性质之后就好做了，因为左边界的取值范围也是不断右移的，所以我们只需要维护两个指针 l 和 r 就行了，一个保存取值范围的最小值，一个保存最大值。然后每次对于一个 j ，符合题意的子区间数量就是 r - l + 1 。而 j 右移一个数之后， l 需要右移，直到 [l, j] 中正好有 K 个不同整数， r 也继续右移，直到[r + 1, j] 中正好有 K - 1 个不同整数。

因为 l 和 r 最多只会移动 n 次，而 j 也只移动了 n 次，所以总体时间复杂度降到了 $O(n)$ 。

代码

c++

        class Solution {
public:
    int subarraysWithKDistinct(vector<int>& A, int K) {
        int n = A.size();
        int cl[n+1], cr[n+1], l = 0, r = 0;
        int res = 0, nl = 0, nr = 0;
        memset(cl, 0, sizeof cl);
        memset(cr, 0, sizeof cr);

        for (int i = 0; i < n; ++i) {
            if (cl[A[i]]++ == 0) nl++;
            while (nl > K) {
                if (--cl[A[l++]] == 0) nl--;
            }
            
            if (cr[A[i]]++ == 0) nr++;
            while (nr >= K) {
                if (--cr[A[r++]] == 0) nr--;
            }
            res += r - l;
        }
        return res;
    }
};

python

        class Solution:
    def subarraysWithKDistinct(self, A: List[int], K: int) -> int:
        n = len(A)
        cl = [0] * (n+1)
        cr = [0] * (n+1)
        l = r = nl = nr = res = 0
        for i in range(n):
            if cl[A[i]] == 0:
                nl += 1
            cl[A[i]] += 1
            while nl > K:
                cl[A[l]] -= 1
                if cl[A[l]] == 0:
                    nl -= 1
                l += 1
            if cr[A[i]] == 0:
                nr += 1
            cr[A[i]] += 1
            while nr >= K:
                cr[A[r]] -= 1
                if cr[A[r]] == 0:
                    nr -= 1
                r += 1
            res += r - l
        return res