Topics › Sliding Window and Two Pointers › Hard Problems

Subarrays with K Different Integers

Medium TUF ID 988 A2Z

LeetCode ↗ Java source ↗ Edit this article ↗

Jump to code: Brute Better Optimal

Problem

Given an integer array nums and an integer k, return the number of good subarrays. A subarray (contiguous, non-empty) is good if it contains exactly k different integers. Subarrays at different positions are counted separately.

Constraints: 1 ≤ n ≤ 2 × 10^4, 1 ≤ nums[i] ≤ n, 1 ≤ k ≤ n. The answer is at most n(n + 1) / 2 ≈ 2 × 10^8, which fits in an int.

Example 1

nums = [1, 2, 1, 2, 3], k = 2
answer = 7          // [1,2], [2,1], [1,2], [2,3], [1,2,1], [2,1,2], [1,2,1,2]

Example 2

nums = [1, 2, 1, 3, 4], k = 3
answer = 3          // [1,2,1,3], [2,1,3], [1,3,4]

Approach 1: Brute force

Intuition. Enumerate subarrays by start i and end j, keeping a set of the values in nums[i..j]. Each time the set size equals k, count the subarray. The number of distinct values can only grow as j moves right, so once it exceeds k, no later end for this start can be good; break.

Algorithm.

1. For each i: create an empty set.
2. For j = i ..: add nums[j]. If the size is k, increment count; if it is greater than k, break.
3. Return count.

static int bruteForce(int[] nums, int k) {
    int n = nums.length, count = 0;
    for (int i = 0; i < n; i++) {
        Set<Integer> seen = new HashSet<>();
        for (int j = i; j < n; j++) {
            seen.add(nums[j]);
            if (seen.size() == k) count++;
            else if (seen.size() > k) break;       // distinct count never drops as j grows
        }
    }
    return count;
}

Complexity. Time O(n^2): with many repeated values the break comes late. Space O(n) for the set in the worst case.

Approach 2: Optimal, “at most k” minus “at most k − 1”

Why a plain window fails. It is tempting to slide a window and count it whenever it has exactly k distinct values. But for a fixed right end, the good subarrays are not a single window; they are a whole range of starts. In [1, 2, 1, 2, 3] with k = 2, the right end at index 3 (value 2) forms a good subarray with each of the starts 0, 1 and 2. A single left pointer can only represent one of them, and if you shrink to the shortest good window you lose the others.

The identity. Let atMost(k) be the number of subarrays with at most k distinct values. Every subarray with at most k distinct values has either exactly k or at most k − 1. Therefore

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

Counting “at most k” with a window. “At most k distinct” is monotonic: if nums[l..r] qualifies, so does every sub-window. So for each right end r, shrink l until the window has at most k distinct values. Then all starts l, l + 1, ..., r work, which contributes r - l + 1 subarrays ending at r. Summing over r counts each qualifying subarray exactly once. A HashMap from value to its count in the window tells us the number of distinct values via map.size(), as long as a key is removed when its count drops to zero.

Algorithm.

1. atMostDistinct(k): return 0 if k < 0. Otherwise for each r, add nums[r] to the map; while the map has more than k keys, remove one occurrence of nums[l] and l++; add r - l + 1.
2. Return atMostDistinct(k) - atMostDistinct(k - 1).

static int optimal(int[] nums, int k) {
    return atMostDistinct(nums, k) - atMostDistinct(nums, k - 1);
}

static int atMostDistinct(int[] nums, int k) {
    if (k < 0) return 0;
    Map<Integer, Integer> freq = new HashMap<>();
    int l = 0, count = 0;
    for (int r = 0; r < nums.length; r++) {
        freq.merge(nums[r], 1, Integer::sum);
        while (freq.size() > k) {
            int x = nums[l];
            if (freq.merge(x, -1, Integer::sum) == 0) freq.remove(x);
            l++;
        }
        count += r - l + 1;                        // starts l..r all give at most k distinct values
    }
    return count;
}

Dry run on nums = [1, 2, 1, 2, 3], k = 2:

12 − 5 = 7. Look at row r = 3: atMost(2) says four subarrays end there with at most two values, atMost(1) says one of them has only one value, so exactly three are good, matching the three starts we found by hand.

Complexity. Time O(n): two window passes, each moving both pointers at most n times, with O(1) average map operations. Space O(n) for the map in the worst case (at most k + 1 keys at any time, so really O(k)).

Because 1 ≤ nums[i] ≤ n, the map can be replaced with an int[n + 1] count array plus a distinct counter. That is the fastest version in Java, since it avoids boxing; the map version is shown because it works for any values.

Edge cases and pitfalls

• k larger than the number of different values. [1, 2, 3] with k = 4 gives 0: both “at most” counts equal the total number of subarrays and cancel out.
• All values equal. [1, 1, 1] with k = 1 gives 6, every subarray.
• The k − 1 < 0 call. atMostDistinct(nums, -1) must return 0 immediately. Without the guard, while (freq.size() > -1) is always true and l runs off the end of the array.
• Removing keys. Remove a value from the map only when its count reaches zero; otherwise freq.size() stops being the number of distinct values in the window.
• Comparing Integer results. freq.merge(x, -1, Integer::sum) == 0 is safe because one side is an int literal, which forces unboxing.

Related problems

• Longest Substring With At Most K Distinct Characters uses the same “at most k distinct” window, but measures length instead of counting.
• Count Number of Substrings is this problem on strings: substrings with exactly k distinct characters.
• Binary Subarrays With Sum uses the same “exactly = at most − at most” identity on sums.
• Count number of Nice subarrays applies the identity to the count of odd numbers.