Subarrays with K Different Integers

Counting questions look like windows but break the usual template. "Exactly k distinct" has no clean shrink rule: adding an element can push the count from k to k + 1, and removing one can pull it from k to k − 1. The fix is a subtraction that turns one hard question into two easy ones.

The problem

Given an array of integers and a number k, count the contiguous subarrays that contain exactly k different values.

• nums = [3, 1, 3, 1, 2], k = 2 → 7. They are [3,1], [3,1,3], [3,1,3,1], [1,3], [1,3,1], [3,1] (starting at index 2) and [1,2].
• nums = [1, 2, 3], k = 3 → 1. Only the whole array has three different values.

Constraints: 1 ≤ n ≤ 2 × 10⁴, 1 ≤ k ≤ n, values between 1 and n.

Clarifying questions

• Do equal subarrays at different positions count separately? Yes, as in the example, where [3,1] counts twice.
• Can k be larger than the number of distinct values? Yes; then the answer is 0.
• Could the answer overflow? There are at most n(n + 1)/2 ≈ 2 × 10⁸ subarrays. Python integers do not overflow; in Java or C++ use a 64-bit type.

Approach 1: every start, growing to the right

For each start, grow the end with a set of values seen. Count each end where the set has exactly k values; stop once it passes k, because it can only grow.

1def k_distinct_brute(nums: list[int], k: int) -> int:2    """Every start, grow right until more than k distinct: O(n^2)."""3    total = 04    for start in range(len(nums)):5        seen: set[int] = set()6        for end in range(start, len(nums)):7            seen.add(nums[end])8            if len(seen) == k:9                total += 110            elif len(seen) > k:11                break12    return total

Time: O(n²) in the worst case (for example when all values are equal and k = 1). At n = 2 × 10⁴ that is 2 × 10⁸ steps: slow in Python, and interviewers expect O(n).

The key insight

Two ideas combine.

First, "at most k" is easy to count with a window. Run the longest-window template with the rule "at most k distinct values". After the shrink, the window start..end is the longest valid window ending at end. Every shorter window ending at end is valid too, because removing elements from the left cannot add distinct values. There are end − start + 1 such windows: starting at start, start + 1, …, end. So add end − start + 1 at each step, and you have counted every valid subarray exactly once, grouped by where it ends.

Second, "exactly k" is a difference. Windows with at most k distinct values are those with exactly k plus those with at most k − 1:

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

Everything with fewer than k distinct values is counted in both terms and cancels. What remains is exactly the subarrays with k.

Why not count "exactly k" with one window directly? Look at [3, 1, 3, 1, 2] at end = 3. The windows ending there with exactly two values start at 0, 1 and 2 ([3,1,3,1], [1,3,1], [3,1]), but the one starting at 3, [1], has only one value. So the valid starts form a range with a left edge and a right edge, and a single window only tracks the left edge. The at-most-k window for k finds the left edge; the at-most-k window for k − 1 finds the right edge. Subtracting the two counts is the same as measuring the width of that range at every step.

Approach 2: two at-most-k windows

1def at_most_k_distinct(nums: list[int], k: int) -> int:2    """Number of subarrays with at most k distinct values."""3    counts: dict[int, int] = {}4    start = 05    total = 06    for end, value in enumerate(nums):7        counts[value] = counts.get(value, 0) + 18        while len(counts) > k:9            leaving = nums[start]10            counts[leaving] -= 111            if counts[leaving] == 0:12                del counts[leaving]13            start += 114        total += end - start + 1      # every window ending at end that starts at or after start15    return total161718def subarrays_with_k_distinct(nums: list[int], k: int) -> int:19    """Exactly k distinct = at most k minus at most k - 1."""20    return at_most_k_distinct(nums, k) - at_most_k_distinct(nums, k - 1)

Dry run of at_most_k_distinct([3, 1, 3, 1, 2], 2):

At end = 4 the value 2 makes three distinct values, so the shrink removes 3, 1 and 3 until only [1, 2] is left.

With k = 1, every value differs from its neighbour, so the window is always one element and the total is 1 + 1 + 1 + 1 + 1 = 5. The answer is 12 − 5 = 7, matching the list in the problem.

Time: O(n). Each call is one window pass where every index enters once and leaves at most once; two calls are still O(n). Space: O(k) for the count map, which never holds more than k + 1 keys.

Edge cases

• k larger than the number of distinct values. atMost(k) and atMost(k − 1) both count every subarray, so the difference is 0.
• k = 1. atMost(0) is 0: the shrink empties the window every step and adds end − start + 1 = 0. The formula then returns the number of subarrays made of one repeated value.
• All values equal. Every subarray has one distinct value. For k = 1 the answer is n(n + 1)/2; for k ≥ 2 it is 0.

Follow-ups

• "Count subarrays with exactly k odd numbers." Same trick: at-most-k with the rule "odd count is at most k", then subtract (Count Number of Nice Subarrays).
• "Count binary subarrays with sum exactly S." With 0s and 1s the sum only grows, so at-most-S minus at-most-(S − 1) works. With arbitrary integers it does not; use prefix sums with a hash map.
• "Longest subarray with at most two distinct values." One at-most-k window, recording the length instead of counting (Fruit Into Baskets).