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. 1 <= A.length <= 20000
2. 1 <= A[i] <= A.length
3. 1 <= K <= A.length

 

双循环滑动窗口

TLE 40/55

class Solution:
    def subarraysWithKDistinct(self, A: List[int], K: int) -> int:
        curLen=0    
        res=0
        for i in range(len(A)):
            cnt=0
            for j in range(i,len(A)):
                curLen=len(set(A[i:j+1]))
                if curLen==K:
                    res+=1
                elif curLen>K:
                    break
        return res

运用哈希表，双指针继续莽

wa

class Solution:
    def subarraysWithKDistinct(self, A: List[int], K: int) -> int:  
        res=0
        i=j=0
        cnt=collections.defaultdict(int)
        while True:
            while j<len(A) and len(cnt)<=K:
                cnt[A[j]]+=1
                j+=1
                if len(cnt)==K:res+=1
            while i<len(A) and len(cnt)>K:
                cnt[A[i]]-=1
                if cnt[A[i]]==0:del cnt[A[i]]
                i+=1
                if len(cnt)==K:res+=1
            if i>=len(A) or j>=len(A):break
        while i<len(A):
            cnt[A[i]]-=1
            if cnt[A[i]]==0:del cnt[A[i]]
            i+=1
            if len(cnt)==K:res+=1
        return res
        

其实移动双指针的时候，可以计算不超过K的，then K=

cnt_at_most(K)-cnt_at_most(K-1)

class Solution:
    def subarraysWithKDistinct(self, A: List[int], K: int) -> int:
        def cnt_at_most(K):
            res=0
            i=j=0
            cnt=collections.defaultdict(int)
            while True:
                #在子数组中不同整数的个数<=K的情况下j指针向右扩张
                while j<len(A) and len(cnt)<=K:
                    #对A[j]加1
                    cnt[A[j]]+=1
                    #移动指针
                    j+=1
                    #满足cnt_at_most(K),res加上j-i
                    if len(cnt)<=K:res+=j-i
                #在子数组中不同整数的个数>K的情况下i指针向右扩张
                while i<len(A) and len(cnt)>K:
                    #对A[i]减1
                    cnt[A[i]]-=1
                    #cnt为0时删除
                    if cnt[A[i]]==0:del cnt[A[i]]
                    #移动指针
                    i+=1
                    #满足cnt_at_most(K),res加上j-i
                    if len(cnt)<=K:res+=j-i
                if j>=len(A):break
            return res
        #cnt_at_most(K)-cnt_at_most(K-1)即为答案
        return cnt_at_most(K)-cnt_at_most(K-1)

 

 

我们再回到第一个解法

如何降维？

经过一番探索发现Google，可以写出这样的解法

class Solution:
    def subarraysWithKDistinct(self, A: List[int], K: int) -> int:
        cnt = {}
        res = i = diffNum = leftForward = 0
        for j in range(len(A)):
            if A[j] not in cnt:
                diffNum += 1
                cnt[A[j]] = 1
            else:
                cnt[A[j]] +=1
            if diffNum == K:
                #当A[i-1] != A[j] and i > 0 的时候置 leftForward 为 0
                if A[i-1] != A[j] and i > 0:
                    leftForward = 0
                #在子数组中不同整数的个数恰好为 K 的情况下i指针向右移动
                while diffNum == K:
                    #注意cnt[A[i]] == 1的情况
                    if cnt[A[i]] == 1:
                        diffNum -= 1
                        del cnt[A[i]]
                    else:
                        cnt[A[i]] -= 1
                    i += 1
                    leftForward += 1
            res += leftForward
        return res