Subarrays with K Different Integers

Subarrays with K Different Integers is the problem that teaches a beautiful counting trick: when "exactly K" is hard, count "at most K" twice and subtract. It turns a tricky window into a routine one.

The slow way first

The obvious idea: enumerate every subarray, drop each into a set, and check whether the set size is exactly k. There are O(n²) subarrays and building the set costs more on top, so this is O(n²) or worse — far too slow for a large array.

The natural instinct is a sliding window. But "exactly k distinct" does not slide cleanly: when you add an element the distinct count can jump, and there is no single window length that captures every valid subarray ending at a position. We need a cleaner quantity to count.

The idea: count "at most", then subtract

The fix is a classic identity:

exactly(k) = atMost(k) − atMost(k − 1)

A subarray with at most k distinct values, minus those with at most k − 1, leaves exactly those with exactly k. And atMost(k) does slide nicely: grow the window with hi; whenever the distinct count exceeds k, shrink from lo until it is valid again. At every hi, the number of valid subarrays ending there is the window length hi − lo + 1.

The key insight: counting hi − lo + 1 at each step adds every subarray that ends at hi and starts anywhere in the current window. Summed over all hi, that totals every valid subarray exactly once.

Walk through it

Step through the animation for atMost(2) on [1, 2, 1, 3]. The hi pointer grows the window; lo only moves when the distinct count exceeds 2. After each valid step we add the window length to count. The window stays valid through [1, 2, 1], then 3 pushes distinct to 3 and lo slides until we are back to [1, 3]. The running total reaches atMost(2) = 8. Doing the same with k = 1 gives atMost(1) = 4, so the answer is 8 − 4 = 4.

Pseudocode

function atMost(nums, k):
    count = 0, lo = 0, freq = empty map
    for hi from 0 to n-1:
        add nums[hi] to freq
        while freq has more than k distinct keys:
            remove nums[lo] from freq (delete key if its count hits 0)
            lo += 1
        count += hi - lo + 1        # subarrays ending at hi
    return count
 
answer = atMost(nums, k) - atMost(nums, k - 1)

The Python solution

def at_most(nums, k):
    count, lo, freq = 0, 0, {}
    for hi, x in enumerate(nums):
        freq[x] = freq.get(x, 0) + 1
        while len(freq) > k:
            left = nums[lo]
            freq[left] -= 1
            if freq[left] == 0: del freq[left]
            lo += 1
        count += hi - lo + 1
    return count
 
def subarrays_k_distinct(nums, k):
    return at_most(nums, k) - at_most(nums, k - 1)

• freq is a dictionary mapping each value in the window to how many times it appears. len(freq) is the distinct count.
• We grow the window by adding nums[hi] to freq.
• The while shrinks from lo whenever distinct exceeds k. We decrement the leftmost value and delete the key when its count hits 0 — that is what actually drops the distinct count.
• Line 10 is the heart of the count: hi − lo + 1 is the number of valid subarrays ending at hi.
• Line 14 applies the identity: the two at_most calls share one helper, so we write the window logic once.

Complexity

Each at_most call is O(n) because hi and lo each advance at most n times total, and we call it twice. The map holds at most k + 1 distinct keys, so extra space is O(k).

When this pattern shows up

Practice