import java.util.HashMap;
class Solution {
public int subarraysWithKDistinct(int[] nums, int k) {
/*
* atMostKDistinct(nums, k):
* Counts the number of subarrays with at most k distinct integers.
* This includes subarrays with 0, 1, 2, ..., up to k distinct integers.
*
* atMostKDistinct(nums, k-1):
* Counts the number of subarrays with at most k-1 distinct integers.
* This includes subarrays with 0, 1, 2, ..., up to k-1 distinct integers.
----------------------------------------------------------------------------------------------------
* exactlyKDistinct(nums, k) = atMostKDistinct(nums, k) - atMostKDistinct(nums, k-1)
*/
return atMostKDistinct(nums, k) - atMostKDistinct(nums, k - 1);
}
private int atMostKDistinct(int[] nums, int k) {
HashMap<Integer, Integer> fMap = new HashMap<>();
int start = 0;
int end = 0;
int count = 0;
/*
* The idea is to keep expanding the right boundary of the window till the count
* of distinct elements in the window is less than or equal to K and
* when the count of distinct elements inside the window becomes more than K,
* start shrinking the window from the left till the count becomes less than or equal
* to K.
* Also for every expansion, keep counting the subarrays as `end – start + 1`
* where right and left are the boundaries of the current window.
*/
while (end < nums.length) {
// Add the current number to the frequency map
fMap.put(nums[end], fMap.getOrDefault(nums[end], 0) + 1);
// If the map size exceeds k, we need to shrink the window
while (fMap.size() > k) {
fMap.put(nums[start], fMap.get(nums[start]) - 1);
if (fMap.get(nums[start]) == 0) {
fMap.remove(nums[start]);
}
start++;
}
// The number of subarrays with at most k distinct elements
// ending at the current 'end' index is (end - start + 1)
if (fMap.size() <= k) { // btw we don't need this condition, it's just for clarity
int windowSize = end - start + 1;
count += windowSize;
}
// Move the end pointer to the right
end++;
}
return count;
}
}