Learning On the Way

992. Subarrays with K Different Integers

Given an array A of positive integers, call a (contiguous, not necessarily distinct) subarray of A good if the number of different integers in that subarray is exactly K. (For example, [1,2,3,1,2] has 3 different integers: 1, 2, and 3.) Return the number of good subarrays of A.

from collections import Counter
class Solution:
    # This solution requires "atMost". If the usual sliding windows method is used, you will find that the "left" pointer will move in both directions, which is not the idea for sliding windows
    def subarraysWithKDistinct(self, A: List[int], K: int) -> int:
        def atMost(K):
            ct = Counter()
            countDiff = 0
            res = i = 0
            for j in range(len(A)):
                ct[A[j]] += 1
                if ct[A[j]] == 1:
                    countDiff += 1

                while countDiff > K:
                    ct[A[i]] -= 1
                    if ct[A[i]] == 0:
                        countDiff -= 1
                    i += 1
                res += j - i + 1
            
            return res
        
        return atMost(K) - atMost(K-1)