LeetCode 992. Subarrays with K Different Integers

Posted on Edited on In 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
class Solution {
public:
    int subarraysWithKDistinct(vector<int>& nums, int k) {
        int n = nums.size();
        // get the count of subarrays that has at most k distinct integers.
        function<int(int)> atMostDistinct = [&](int k) {
            int res = 0;
            vector<int> freq(n + 1); // given "1 <= nums[i] <= nums.length".
            // cnt is the count of distinct integers within [left, right].
            for (int left = 0, right = 0, cnt = 0; right < n; ++right) {
                if (++freq[nums[right]] == 1) ++cnt;
                while (cnt > k)
                    if (--freq[nums[left++]] == 0) --cnt;
                // add the count of new subarrays that ends with nums[right].
                res += right - left + 1;
            }
            return res;
        };
        return atMostDistinct(k) - atMostDistinct(k - 1);
    }
};