Subarrays with K Different Integers [Hard] — atMost(K)−atMost(K−1)

Problem Statement

Given integer array nums and integer k, return the number of good subarrays (exactly k different integers).

Input: nums=[1,2,1,2,3], k=2 → Output: 7

Intuition

exactly(k) = atMost(k) - atMost(k-1).

atMost(k): track a window with at most k distinct values. For each valid right, add right-left+1 subarrays.

Solutions

C++

int atMost(vector<int>& A, int k){
    unordered_map<int,int> cnt; int l=0,res=0;
    for(int r=0;r<A.size();r++){
        cnt[A[r]]++;
        while(cnt.size()>k){if(--cnt[A[l]]==0)cnt.erase(A[l]);l++;}
        res+=r-l+1;
    }
    return res;
}
int subarraysWithKDistinct(vector<int>& nums, int k){return atMost(nums,k)-atMost(nums,k-1);}

Java

public int subarraysWithKDistinct(int[] nums, int k){return atMost(nums,k)-atMost(nums,k-1);}
int atMost(int[] A, int k){
    Map<Integer,Integer> cnt=new HashMap<>();int l=0,res=0;
    for(int r=0;r<A.length;r++){
        cnt.merge(A[r],1,Integer::sum);
        while(cnt.size()>k){cnt.merge(A[l],-1,Integer::sum);if(cnt.get(A[l])==0)cnt.remove(A[l]);l++;}
        res+=r-l+1;
    }
    return res;
}

Python

from collections import defaultdict

def subarraysWithKDistinct(nums: list[int], k: int) -> int:
    def at_most(k):
        cnt = defaultdict(int)
        left = res = 0
        for right, v in enumerate(nums):
            cnt[v] += 1
            while len(cnt) > k:
                cnt[nums[left]] -= 1
                if cnt[nums[left]] == 0: del cnt[nums[left]]
                left += 1
            res += right - left + 1
        return res
    return at_most(k) - at_most(k-1)

Complexity

Time: O(n), Space: O(k)