Leetcode 992: Subarrays with K Different Integers

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

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

Notes:

This question looks like a two-pointer sliding window problem. Left is the starting position. Once we reach the state of having K different integers, we need to know the last position of the first kind of number, which is label as i. Then there are (i - Left) number of string.

Now we can add the next number into account.

(1) If this number shows up previously, then we just need to update the last appearing position of it. The added number of string is still (i - Left);

(2) If this number is new element, then we need to remove the first kind of number (or the one with the smallest last position, newLeft). update Left to newLeft. update memorization, add (i - newLeft) into the final res

See the code below:

class Solution {
public:
    int subarraysWithKDistinct(vector<int>& A, int K) {
        int res = 0, left = -1, i = 0;
        unordered_map<int, int> mp;//<num, indx>
        set<pair<int, int>> st;//<<indx, num>>
        for(; i<A.size(); ++i) {
            if(mp.size() == K) break;
            if(mp.count(A[i])) st.erase({mp[A[i]], A[i]});
            mp[A[i]] = i;
            st.insert({i, A[i]});
        }
        if(mp.size() < K) return res;
        res += st.begin()->first - left;
        for(; i<A.size(); ++i) {
            if(mp.count(A[i])) {
                st.erase({mp[A[i]], A[i]});
            }
            else {
                left = st.begin()->first;
                mp.erase(st.begin()->second);
                st.erase(st.begin());
            }
            mp[A[i]] = i;
            st.insert({i, A[i]});
            res += st.begin()->first - left;
        }
        return res;
    }
};

There is much clever solution found online, here is the link.

The key is using the results from a simple question.

Find the code below:

class solution {
public:
    int subarraysWithKDistinct(vector<int>& A, int K) {
        return atMostK(A, K) - atMostK(A, K - 1);
    }
private:
    int atMostK(vector<int>& A, int K) {
        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++;
                i++;
            }
            res += j - i + 1;
        }
        return res;
    }
};