Subarrays with K Different Integers

Given an array A of positive integers, call a (contiguous, not necessarily distinct) subarray of A good if the number of different integers in that subarray is exactly K.

(For example, [1,2,3,1,2] has 3 different integers: 1, 2, and 3.)

Return the number of good subarrays of A.

Example 1:

Input: A = [1,2,1,2,3], K = 2
Output: 7
Explanation: Subarrays formed with exactly 2 different integers: [1,2], [2,1], [1,2], [2,3], [1,2,1], [2,1,2], [1,2,1,2].

Example 2:

Input: A = [1,2,1,3,4], K = 3
Output: 3
Explanation: Subarrays formed with exactly 3 different integers: [1,2,1,3], [2,1,3], [1,3,4].

Note:

  1. 1 <= A.length <= 20000

  2. 1 <= A[i] <= A.length

  3. 1 <= K <= A.length

class Solution {
    public int subarraysWithKDistinct(int[] A, int K) {
        return numberOfSubArraysWithAtMostKElements(A, K) - numberOfSubArraysWithAtMostKElements(A, K - 1);
    }

    public int numberOfSubArraysWithAtMostKElements(int[] A, int K) {
        HashMap<Integer, Integer> map = new HashMap<>();
        int start = 0, unique = 0;
        int count = 0;
        for (int i = 0; i < A.length; i++) {
            if (!map.containsKey(A[i])) {
                unique++;
                map.put(A[i], 1);
            } else {
                if (map.get(A[i]) == 0)
                    unique++;
                map.put(A[i], map.get(A[i]) + 1);
            }
            while (unique > K) {
                map.put(A[start], map.get(A[start]) - 1);
                if (map.get(A[start]) == 0)
                    unique--;
                start++;
            }
            // We are counting all the subarrays between i and start
            // which will end at i.
            count += i - start + 1;
        }
        return count;
    }
}

Last updated