[leetcode] 992. K个不同整数的子数组

题目

给定一个正整数数组 A，如果 A 的某个子数组中不同整数的个数恰好为 K，则称 A 的这个连续、不一定独立的子数组为好子数组。

（例如，[1,2,3,1,2] 中有 3 个不同的整数：1，2，以及 3。）

返回 A 中好子数组的数目。

 

示例 1：

输出：A = [1,2,1,2,3], K = 2
输入：7
解释：恰好由 2 个不同整数组成的子数组：[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
解释：恰好由 3 个不同整数组成的子数组：[1,2,1,3], [2,1,3], [1,3,4].

 

提示：

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

解题思路

这道题的暴力解法是O(n^2)，并且蛮明显是用滑动窗口做的。

但是发现被标为hard所以估计要找O(n)的算法。

滑动窗口有一个很明显的特征是不会反复，就是左指针left和右指针right都只会增长，不会减小，这样可以达到O(n)。

那么对于每个A[right]，考虑两种情况，一种是 A[right] 是已经出现过的数字，那么只要以A[right-1]为结尾的“好子数组”都肯定满足“好子数组”的条件，这个时候只需要记录以A[right-1]为结尾的“好子数组”的数目，然后加上left向右收缩得到的增量；一种是 A[right] 是新的数字，一般情况下 [left, right-1]达到了K， 那么当前数组[left, right]的不同数字数目应该达到了K+1，这个时候 left 的位置不符合条件，需要向右收缩。

综上，left和right都只需要增长不需要回溯，所以复杂度是O(n)。

另外一个难点是处理left，right和当前数字的细节。我处理得蛮生硬的，然后因为细节处理错了好几次，所以先贴我的代码，再贴一个我觉得细节处理得很清晰的123场周赛中国区第一@LongChen的代码。思路是一样的。

我的代码：

class Solution {
public:
    int subarraysWithKDistinct(vector& A, int K) {
        map freq;
        for(int i = 0; i < A.size(); i++){
            int val = A[i];
            freq[val] = 0;
        }
        
        int res = 0;
        int l = 0, r = 0, diff = 0;
        int cur = 1, sum = 0;
        
        while(r < A.size()){
            while(r < A.size() && diff < K){
                if(freq[A[r]] == 0){
                    diff += 1;
                }
                freq[A[r]] += 1;
                if(diff == K){
                    break;
                }
                r += 1;
            }
            if(r == A.size() && diff < K){
                return 0;
            }
            
            while(l <= r - K + 1 && diff > K){
                if(freq[A[l]] == 1){
                    diff -= 1;
                }
                freq[A[l]] -= 1;
                l += 1;
            }

            int begin_l = l;
            while(l <= r - K + 1 && diff == K){
                if(l > begin_l){
                    cur += 1;
                }
                if(freq[A[l]] == 1){
                    break;
                }
                freq[A[l]] -= 1;
                l += 1;
            }
            sum += cur;

            r += 1;
            if(r < A.size()){
                if(freq[A[r]] == 0){
                    diff += 1;
                    cur = 1;
                }
                freq[A[r]] += 1;
            }
            
        }
        return sum;
    }
};

@LongChen的代码：

class Solution
{
public:
    int subarraysWithKDistinct(vector& A, int K)
    {
        int size = A.size();
        int hash[size+1] = {};
        int cnt = 0;
        int ans = 0;
        int l = 0, r = -1;
        int step = 1;
        while(++r < size)
        {
            if(++hash[A[r]] == 1)
                cnt++;
            if(cnt < K)
                continue;
            if(cnt > K)
            {
                while(cnt > K)
                {
                    if(--hash[A[l++]] == 0)
                        cnt--;
                }
                step = 1;
            }
            if(cnt == K)
            {
                while(hash[A[l]] > 1)
                    --hash[A[l++]], step++;
                ans += step;
            }
        }
        return ans;
    }
};