992. Subarrays with K Different Integers

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. 1 <= A.length <= 20000

2. 1 <= A[i] <= A.length

3. 1 <= K <= A.length

// Sliding Window
int subarraysWithAtMostKDistinct(vector<int>& A, int K) { // time: O(n); space: O(n)
    int i = 0, res = 0;
    unordered_map<int, int> count;
    for (int j = 0; j < A.size(); ++j) {
        if (!count[A[j]]++) --K;
        while (K < 0) {
            if (!--count[A[i++]]) ++K;
        }
        res += j - i + 1; // the total number of subarrays ending at j that contain at most K distinct
    }
    return res;
}
int subarraysWithKDistinct(vector<int>& A, int K) {
    return subarraysWithAtMostKDistinct(A, K) - subarraysWithAtMostKDistinct(A, K - 1);
}

// Sliding Window
int atMost(vector<int>& A, int K) {
    int start = 0, end = 0, counter = K, res = 0;
    unordered_map<int, int> record;
    while (end < A.size()) {
        if (record[A[end++]]++ == 0) --counter;
        while (counter < 0) {
            if (record[A[start++]]-- == 1) ++counter;
        }
        res += end - start; // the total number of subarrays ending at end that contain at most K distinct
    }
    return res;
}
int subarraysWithKDistinct(vector<int>& A, int K) { // time: O(n); space: O(n)
    return atMost(A, K) - atMost(A, K - 1);
}

// Sliding Window
int subarraysWithKDistinct(vector<int>& A, int K) { // time: O(n); space: O(n)
    vector<int> m(A.size() + 1, 0);
    int res = 0;
    for (int i = 0, j = 0, prefix = 0, cnt = 0; i < A.size(); ++i) {
        if (m[A[i]]++ == 0) ++cnt;
        // reset
        if (cnt > K) {
            --m[A[j++]];
            --cnt;
            prefix = 0;
        }
        // only keep track of the number of subarrays with the same number of distinct integers
        while (m[A[j]] > 1) {
            ++prefix;
            --m[A[j++]];
        }
        if (cnt == K) res += prefix + 1;
    }
    return res;
}