题目
给定一个正整数数组 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时,左指针向右移动。当右指针走到头,子数组也<K时,结束整体循环。时间复杂度 o ( n 2 ) o(n^2) o(n2),这个方法超时了
双指针法:看了题解才想出来。。。主要思想是:包含K个整数的子数组个数,应该是包含0~K个整数的子数组个数减去包含0~K-1个整数的子数组个数。
求包含0~K个不同整数的子数组数量,其实和713. Subarray Product Less Than 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
代码
class Solution:
def subarraysWithKDistinct(self, nums: List[int], k: int) -> int:
# 1: 0~k unique numbers, 2: 0~k-1 unique numbers
left1, left2 = 0, 0
res = 0
num_cnt1, num_cnt2 = {}, {}
for right in range(len(nums)):
num_cnt1[nums[right]] = num_cnt1.get(nums[right], 0) + 1
num_cnt2[nums[right]] = num_cnt2.get(nums[right], 0) + 1
while len(num_cnt1) > k:
num_cnt1[nums[left1]] -= 1
if num_cnt1[nums[left1]] == 0:
num_cnt1.pop(nums[left1])
left1 += 1
while len(num_cnt2) > k - 1:
num_cnt2[nums[left2]] -= 1
if num_cnt2[nums[left2]] == 0:
num_cnt2.pop(nums[left2])
left2 += 1
if len(num_cnt1) == k and len(num_cnt2) == k - 1:
res += left2 - left1
return res