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;
}
};