Leetcode 992 Subarrays with K Different Integers

发表于 |

给定一个正整数数组A，现在定义A的一个子串是good当且仅当这个子串内包含恰好k个不同的数！
例如，[1,2,3,1,2]有3个不同的整数1,2,3
返回A中good的子串的个数。

示例 1:
输入: A = [1,2,1,2,3], K = 2
输出: 7
解释: 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].

示例 2:
输入: A = [1,2,1,3,4], K = 3
输出: 3
解释: Subarrays formed with exactly 3 different integers: [1,2,1,3], [2,1,3], [1,3,4].

Note:

  1. 1 <= A.length <= 20000
  2. 1 <= A[i] <= A.length
  3. 1 <= K <= A.length

思路分析：
inspired by @lee215
关键就在于这个思路，定义f(k)表示数组A中至多包含K个不同的数的子串个数
那么题目所求包含恰好k个不同的数的子串个数则可表示为f(k) - f(k-1)

显然，f(k)的计算是非常简单的，双指针思路如下：
对于每个A[j]，计算以A[j]结尾的最多包含k个不同的数的子串个数，然后求和即可。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
# inspired by @lee215
class Solution(object):
    def subarraysWithKDistinct(self, A, K):
        def solve(A, K):
            i = 0
            res = 0
            d = collections.Counter()
            for j in range(len(A)):
                d[A[j]] += 1
                if d[A[j]] == 1: K -= 1
                while K < 0:
                    d[A[i]] -= 1
                    if d[A[i]] == 0: K += 1
                    i += 1
                res += j - i + 1
            return res
        return solve(A, K) - solve(A, K-1)