题目描述
给定一个正整数数组 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 的区间个数。但是这种方法时间复杂度是 ,一定会超时,所以考虑其他方法。
现在考虑右边界为 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 次,所以总体时间复杂度降到了 。
代码
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
后记
其实这题想起来可能好想,但是写起来容易写错,因为区间范围需要好好琢磨。这一类问题统称为“窗口滑动”问题,都是不特别难,想清楚了两个状态之间窗口如何滑动就行了。
作者简介:godweiyang,知乎同名,华东师范大学计算机系硕士在读,方向自然语言处理与深度学习。喜欢与人分享技术与知识,期待与你的进一步交流~