题目描述

给定一个正整数数组 A，称 a（连续的，不一定不同）是 A 的 漂亮的 子数组当其中不同数字的个数正好为 K。
（例如，[1,2,3,1,2] 有三个不同的数字：1，2 和 3。

返回 A 中漂亮子数组的个数。

样例

输入：A = [1,2,1,2,3], K = 2
输出：7
解释：正好由两个不同数字形成的子数组：[1,2]，[2,1]，[1,2]，[2,3]，[1,2,1]，[2,1,2]，[1,2,1,2]。

输入：A = [1,2,1,3,4], K = 3
输出：3
解释：正好由三个不同数字形成的子数组：[1,2,1,3]，[2,1,3]，[1,3,4]。

注意

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

算法

(双指针 / 滑动窗口) $O(n)$

1. 对于每一个位置 i，求出 lower[i] 表示达到正好 K 个不同数字的最小下标，和 upper[i] 表示达到正好 K - 1 个不同数字的最小下标。累计每个位置的 upper[i] - lower[i] 就是答案。
2. 具体求的过程可以用双指针（滑动窗口）来求解。对于求 lower[i] 来说，如果 i > j，则显然有 lower[i] >= lower[j]。对于求 upper[i] 同理。

时间复杂度

• 每个位置仅扫描常数次，故总时间复杂度为 $O(n)$。

空间复杂度

• 需要额外的数组空间，故总空间复杂度为 $O(n)$。

C++ 代码

class Solution {
public:

    void find(vector<int> &last, const vector<int> &A, int K) {
        int n = A.size();
        int j = 0, cur = 0;
        vector<int> vis(n + 1, 0);
        for (int i = 0; i < n; i++) {
            if (vis[A[i]] == 0)
                cur++;
            vis[A[i]]++;
            while (cur == K + 1) {
                if (vis[A[j]] == 1)
                    cur--;
                vis[A[j]]--;
                j++;
            }
            last[i] = j;
        }
    }

    int subarraysWithKDistinct(vector<int>& A, int K) {
        int n = A.size();
        vector<int> lower(n, 0), upper(n, 0);
        find(upper, A, K - 1);
        find(lower, A, K);
        int ans = 0;
        for (int i = 0; i < n; i++)
            ans += upper[i] - lower[i];

        return ans;
    }
};