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)$，这个方法超时了

双指针法：看了题解才想出来。。。主要思想是：包含K个整数的子数组个数，应该是包含0～K个整数的子数组个数减去包含0～K-1个整数的子数组个数。

求包含0～K个不同整数的子数组数量，其实和乘积小于K的子数组很相似，每次都计算窗口内包含最右边元素的子数组个数，个数就是right - left + 1

保存2个窗口，1个窗口计算含有0～K个不同元素的子数组数量，1个窗口计算含有0～K-1个不同元素的子数组数量，两个窗口相减就能求出恰好包含K个不同元素的子数组数量了。所以是right1 - left1 + 1 - (right2 - left2 + 1)，然后发现2个窗口共用1个right，所以简化为：right - left1 + 1 - (right - left2 + 1) = left2 - left1

这道题还有一点需要注意的，因为list转换成set的复杂度是$o(n)$，所以如果不用字典来判断队列中元素的数量，而是用len(set(queue))的话会超时

另外因为每次需要从queue队头元素出队，所以可以用collections.deque.popleft()来优化，此时出队的时间复杂度变成$o(1)$

代码

class Solution:
    def subarraysWithKDistinct(self, A: List[int], K: int) -> int:
        k1_queue, k2_queue = collections.deque(), collections.deque()
        k1_dict, k2_dict = {}, {}
        l1, l2 = 0, 0
        cnt = 0
        for num in A:
            k1_queue.append(num)
            k2_queue.append(num)
            k1_dict[num] = k1_dict.get(num, 0) + 1
            k2_dict[num] = k2_dict.get(num, 0) + 1
            while len(k1_dict) > K:
                pop_item = k1_queue.popleft()
                l1 += 1
                k1_dict[pop_item] -= 1
                if k1_dict[pop_item] == 0:
                    k1_dict.pop(pop_item)
            while len(k2_dict) > K - 1:
                pop_item = k2_queue.popleft()
                l2 += 1
                k2_dict[pop_item] -= 1
                if k2_dict[pop_item] == 0:
                    k2_dict.pop(pop_item)
            cnt += (l2 - l1)
        return cnt