Leetcode-992 Subarrays with K Differant Integers

题目描述

https://leetcode.com/problems/subarrays-with-k-different-integers/

给定一个数组 A，求有 k 个不同元素组成的子数组个数。

例如：

1
2
3

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].

测试用例

根据题意，数组长度 n 的取值区间为 [1, 20000]，数组元素取值区间为 [1, n]，说明最多可能有 n 个不同元素，k 取值区间也是 [1, n]，需要考虑 k > count(distinct A[i]) 的情况。考虑如下测试用例：

// Case 0: k > count(distinct A[i])
A = [1, 1, 1, 1, 1], k = 2
Output: 0
// Case 1: k == count(distinct A[i])
A = [1, 1, 2, 3], k = 3
// Case 2: k < count(distinct A[i])

分析

第一眼看应该是使用滑窗，根据之前的总结滑窗的模板如下

class Solution:
    def slidingWindowTemplate(s, t):
        # 定义其他变量，存储结果，存储判断条件等变量
        begin = end = 0
        
        while f(end): # case 1: slide end
            # do something
            end += 1
            while g(): # case 2: slide begin
                # do something
                 begin += 1
                    
                # meet some condition

使用滑窗需要思考两个问题：

什么时候滑动 end 指针？
什么时候滑动 begin 指针？

经过思考直接求解不好求，我们换个思路：最多包含 k 个不同元素的子数组有多少个？滑窗 + HashMap 是解决这类问题最常用的方案。代码如下：

def atMostKDistinct(A, k):
    begin = end = 0
    window = collections.defaultdict(int)
    n, distinct = len(A), 0 # distinct indicate count of sliding window
    ans = 0
    
    while end < n:
        e = A[end]
        if window[e] == 0:
            distinct += 1
        window[e] += 1
        end += 1
        while distinct > k:
            b = A[begin]
            window[b] -= 1
            if window[b] == 0:
                distinct -= 1
            begin += 1
        # [begin, end) has end - begin items that is numbers of subarray end of A[end]
        ans += end - begin
    return ans

记 f(k) 为小于等于 k 个不同元素的子数组个数，则 f(k-1) 为小于等于 k - 1 个不同元素的子数组个数。所以，确切的有 k 个不同元素的子数组个数 = f(k) - f(k-1)。

代码

根据如上分析，完整代码如下：

class Solution:
    def subarraysWithKDistinct(self, A: List[int], k: int) -> int:
        def atMostKDistinct(k):
            begin = end = 0
            window = collections.defaultdict(int)
            n, distinct = len(A), 0 # distinct indicate count of sliding window
            ans = 0

            while end < n:
                e = A[end]
                if window[e] == 0:
                    distinct += 1
                window[e] += 1
                end += 1
                while distinct > k:
                    b = A[begin]
                    window[b] -= 1
                    if window[b] == 0:
                        distinct -= 1
                    begin += 1
                # [begin, end) has end - begin items that is end - begin subarrays
                ans += end - begin
            return ans
        return atMostKDistinct(k) - atMostKDistinct(k-1)

在 [begin, end) 区间中，以 A[end-1] 结尾的子数组个数是 end - begin，所以 ans += end - begin