'''
keys: If the problem talks about continuous subarrays or substrings, 
the sliding window technique may help solve it in a linear time.
Solutions:
Similar:340
T:
S:
'''
from typing import List

import collections

class Solution1:
    # official
    # O(N) for S and T, N as the length of A
    '''
    This problem will be a very typical sliding window, 
    if it asks the number of subarrays with at most K distinct elements.  
    Just need one more step to reach the folloing equation: 
    exactly(K) = atMost(K) - atMost(K-1)
    A = [1,2,1,2,3], K = 2
    good subarrays: [1,2], [2,1], [1,2], [2,3], [1,2,1], [2,1,2], [1,2,1,2]
    '''
    def subarraysWithKDistinct1(self, A: List[int], K: int) -> int:
        return self.atMostK(A, K) - self.atMostK(A, K - 1)

    def atMostK(self, A, K):
        count = collections.Counter()
        res = i = 0
        for j in range(len(A)):
            if count[A[j]] == 0: 
                K -= 1
            count[A[j]] += 1
            while K < 0:
                count[A[i]] -= 1
                if count[A[i]] == 0: 
                    K += 1
                i += 1
            res += j - i + 1
        return res


# official
class Window:
    def __init__(self):
        self.count = collections.Counter()
        self.nonzero = 0

    def add(self, x):
        self.count[x] += 1
        if self.count[x] == 1:
            self.nonzero += 1

    def remove(self, x):
        self.count[x] -= 1
        if self.count[x] == 0:
            self.nonzero -= 1


class Solution:
    # official
    # O(N) for S and T, N as the length of A
    def subarraysWithKDistinct(self, A: List[int], K: int) -> int:
        window1 = Window()
        window2 = Window()
        ans = left1 = left2 = 0

        for right, x in enumerate(A):
            window1.add(x)
            window2.add(x)

            while window1.nonzero > K:
                window1.remove(A[left1])
                left1 += 1

            while window2.nonzero >= K:
                window2.remove(A[left2])
                left2 += 1

            ans += left2 - left1

        return ans