Subarrays with K Different Integers
Problem
Return the number of subarrays of nums that contain exactly k distinct integers.
nums = [1,2,1,2,3], k = 27def subarraysWithKDistinct(nums, k):
def atMost(K):
cnt = {}
l = total = 0
for r, x in enumerate(nums):
cnt[x] = cnt.get(x, 0) + 1
while len(cnt) > K:
cnt[nums[l]] -= 1
if cnt[nums[l]] == 0: del cnt[nums[l]]
l += 1
total += r - l + 1
return total
return atMost(k) - atMost(k - 1)function subarraysWithKDistinct(nums, k) {
function atMost(K) {
const cnt = new Map();
let l = 0, total = 0;
for (let r = 0; r < nums.length; r++) {
cnt.set(nums[r], (cnt.get(nums[r]) || 0) + 1);
while (cnt.size > K) {
cnt.set(nums[l], cnt.get(nums[l]) - 1);
if (cnt.get(nums[l]) === 0) cnt.delete(nums[l]);
l++;
}
total += r - l + 1;
}
return total;
}
return atMost(k) - atMost(k - 1);
}class Solution {
public int subarraysWithKDistinct(int[] nums, int k) {
return atMost(nums, k) - atMost(nums, k - 1);
}
int atMost(int[] nums, int K) {
Map<Integer, Integer> cnt = new HashMap<>();
int l = 0, total = 0;
for (int r = 0; r < nums.length; r++) {
cnt.merge(nums[r], 1, Integer::sum);
while (cnt.size() > K) {
cnt.merge(nums[l], -1, Integer::sum);
if (cnt.get(nums[l]) == 0) cnt.remove(nums[l]);
l++;
}
total += r - l + 1;
}
return total;
}
}int atMost(vector<int>& nums, int K) {
unordered_map<int, int> cnt;
int l = 0, total = 0;
for (int r = 0; r < (int)nums.size(); r++) {
cnt[nums[r]]++;
while ((int)cnt.size() > K) {
if (--cnt[nums[l]] == 0) cnt.erase(nums[l]);
l++;
}
total += r - l + 1;
}
return total;
}
int subarraysWithKDistinct(vector<int>& nums, int k) { return atMost(nums, k) - atMost(nums, k - 1); }Explanation
Counting subarrays with exactly k distinct values directly is awkward, because a sliding window does not naturally pin a count to an exact number. The clever fix is to count at most instead: exactly(k) = atMost(k) - atMost(k - 1).
The helper atMost(K) counts subarrays with no more than K distinct integers. It slides a window with a count map cnt. As r moves right and adds nums[r], if the window ever holds more than K distinct values, we shrink from the left, removing values and deleting them from the map when their count hits zero.
For each r, every subarray ending at r and starting anywhere from l to r has at most K distinct values, which is r - l + 1 subarrays. We add that to total.
Subtracting atMost(k - 1) from atMost(k) removes all the subarrays with fewer than k distinct values, leaving precisely those with exactly k.
Example: nums = [1,2,1,2,3], k = 2. atMost(2) counts more windows than atMost(1), and their difference is 7, the number of subarrays with exactly two distinct integers.