題目

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.

Example 1:

Input: A = [1,2,1,2,3], K = 2
Output: 7
Explanation: Subarrays formed with exactly 2 different integers: [1,2], [2,1], [1,2], [2,3], [1,2,1], [2,1,2], [1,2,1,2].

Example 2:

Input: A = [1,2,1,3,4], K = 3
Output: 3
Explanation: Subarrays formed with exactly 3 different integers: [1,2,1,3], [2,1,3], [1,3,4].

Note:

• 1 <= A.length <= 20000
• 1 <= A[i] <= A.length
• 1 <= K <= A.length

題目大意

這道題考察的是滑動窗口的問題。

給出一個數組 和 K，K 代表窗口能能包含的不同數字的個數。K = 2 代表窗口內只能有 2 個不同的數字。求數組中滿足條件 K 的窗口個數。

解題思路

如果只是單純的滑動窗口去做，會錯過一些解。比如在例子 1 中，滑動窗口可以得到 [1,2], [1,2,1], [1,2,1,2], [2,1,2], [1,2], [2,3], 會少 [2,1] 這個解，原因是右邊窗口滑動到最右邊了，左邊窗口在縮小的過程中，右邊窗口不會再跟着動了。有同學可能會説，每次左邊窗口移動的時候，右邊窗口都再次從左邊窗口的位置開始重新滑動。這樣做確實可以，但是這樣做完會發現超時。因為中間包含大量的重複計算。

這道題就需要第 3 個指針。原有滑動窗口的 2 個指針，右窗口保留這個窗口裏面最長的子數組，正好有 K 個元素，左窗口右移的邏輯不變。再多用一個指針用來標識正好有 K - 1 個元素的位置。那麼正好有 K 個不同元素的解就等於 ans = atMostK(A, K) - atMostK(A, K - 1)。最多有 K 個元素減去最多有 K - 1 個元素得到的窗口中正好有 K 個元素的解。

以例子 1 為例，先求最多有 K 個元素的窗口個數。

[1]     
[1,2], [2]     
[1,2,1], [2,1], [1]  
[1,2,1,2], [2,1,2], [1,2], [2]  
[2,3], [3]

每當窗口滑動到把 K 消耗為 0 的時候，res = right - left + 1 。為什麼要這麼計算，right - left + 1 代表的含義是，終點為 right，至多為 K 個元素的窗口有多少個。[left,right], [left + 1,right], [left + 2,right] …… [right,right]。這樣算出來的解是包含這道題最終求得的解的，還多出了一部分解。多出來的部分減掉即可，即減掉最多為 K - 1 個元素的解。

最多為 K - 1 個元素的解如下：

[1]
[2]
[1]
[2]
[3]

兩者相減以後得到的結果就是最終結果：

[1,2]    
[1,2,1], [2,1]  
[1,2,1,2], [2,1,2], [1,2]  
[2,3]

參考代碼

package leetcode

func subarraysWithKDistinct(A []int, K int) int {
	return subarraysWithKDistinctSlideWindow(A, K) - subarraysWithKDistinctSlideWindow(A, K-1)
}

func subarraysWithKDistinctSlideWindow(A []int, K int) int {
	left, right, counter, res, freq := 0, 0, K, 0, map[int]int{}
	for right = 0; right < len(A); right++ {
		if freq[A[right]] == 0 {
			counter--
		}
		freq[A[right]]++
		for counter < 0 {
			freq[A[left]]--
			if freq[A[left]] == 0 {
				counter++
			}
			left++
		}
		res += right - left + 1
	}
	return res
}