1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37

// @Title: K 个不同整数的子数组 (Subarrays with K Different Integers)
// @Author: 15816537946@163.com
// @Date: 2021-02-10 09:48:19
// @Runtime: 60 ms
// @Memory: 6.9 MB
func subarraysWithKDistinct(A []int, K int) (ans int) {
	n := len(A)
	num1 := make([]int, n+1)
	num2 := make([]int, n+1)
	var tot1, tot2, left1, left2 int
	for _, v := range A {
		if num1[v] == 0 {
			tot1++
		}
		num1[v]++
		if num2[v] == 0 {
			tot2++
		}
		num2[v]++
		for tot1 > K {
			num1[A[left1]]--
			if num1[A[left1]] == 0 {
				tot1--
			}
			left1++
		}
		for tot2 > K-1 {
			num2[A[left2]]--
			if num2[A[left2]] == 0 {
				tot2--
			}
			left2++
		}
		ans += left2 - left1
	}
	return ans
}