992. Subarrays with K Different Integers

Reference: LeetCode
Difficulty: Hard

Problem

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.

Note:

• The resulting subarrays is not required to be distinct.
• 1 <= A.length <= 20000
• 1 <= A[i] <= A.length
• 11 <= K <= A.length

Example:

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].

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].

Analysis

Brute-Force

For each subarray ($\sim N^2$ in total), check if there are K distinct elements.

public int subarraysWithKDistinct(int[] A, int K) {
  int n = A.length;
  int count = 0;
  for (int i = 0; i < n; ++i) {
    for (int j = i; j < n; ++j) {
      int subLen = j - i + 1;
      if (subLen < K) continue; // length does not respect
      Set<Integer> set = new HashSet<>();
      for (int t = i; t <= j; ++t) { // for each element in subarray
        set.add(A[t]);
      }
      if (set.size() == K) ++count;
    }
  }
  return count;
}

Time: $O(N^3)$
Space: $O(N)$

Hash Set

For each element A[i] as the starting element of a subarray, we can use a hash set to control the current number of elements. Specifically, we first keep adding elements to the set until its size reaches K or there are no more elements that could be added.

public int subarraysWithKDistinct(int[] A, int K) {
  int n = A.length;
  int count = 0;
  for (int i = 0; i < n; ++i) {
    Set<Integer> set = new HashSet<>();
    int j = i;
    while (j < n && set.size() < K) {
      set.add(A[j]);
      ++j;
    } // set's size equals K or j >= n
    if (set.size() == K) {
      ++count;
      while (j < n) { // keep adding to the set
        set.add(A[j]);
        if (set.size() == K) {
          ++count;
          ++j;
        } else {
          break;
        }
      }
    }
  }
  return count;
}

Time: $O(N^2)$
Space: $O(K)$ since the size of the set is limited by $K$.

Sliding Window

The idea is to calculate the number of subarrays with at most K distinct characters and with at most K - 1 distinct characters and then do subtraction.

Example of calculating the count of at most K = 3 in [1, 2, 1, 3, 4]:

Notice why count += end - start + 1 works here.

public int subarraysWithKDistinct(int[] A, int K) {
  int n = A.length;
  return subarraysWithAtMostKDistinct(A, K) - subarraysWithAtMostKDistinct(A, K - 1);
}

private int subarraysWithAtMostKDistinct(int[] A , int K) {
  int n = A.length;
  int count = 0;
  Map<Integer, Integer> map = new HashMap<>(); // count the elements
  for (int start = 0, end = 0; end < n; ++end) {
    map.put(A[end], map.getOrDefault(A[end], 0) + 1); // add one
    // size exceeds
    while (map.size() > K) {
      int c = map.get(A[start]);
      map.put(A[start], c - 1); // minus one
      if (c - 1 == 0) {
        map.remove(A[start]);
      }
      ++start;
    }
    count += end - start + 1;
  }
  return count;
}

Or in one-pass:

public int subarraysWithKDistinct(int[] A, int K) {
  int n = A.length;
  int count = 0;
  Map<Integer, Integer> window1 = new HashMap<>();
  Map<Integer, Integer> window2 = new HashMap<>();
  for (int start1 = 0, start2 = 0, end = 0; end < n; ++end) {
    window1.put(A[end], window1.getOrDefault(A[end], 0) + 1); // add one
    window2.put(A[end], window2.getOrDefault(A[end], 0) + 1); // add one
    // size exceeds
    while (window1.size() > K) {
      int c = window1.get(A[start1]);
      window1.put(A[start1], c - 1);
      if (c - 1 == 0) {
        window1.remove(A[start1]);
      }
      ++start1;
    }
    while (window2.size() > K - 1) {
      int c = window2.get(A[start2]);
      window2.put(A[start2], c - 1);
      if (c - 1 == 0) {
        window2.remove(A[start2]);
      }
      ++start2;
    }
    count += start2 - start1;
  }
  return count;
}

Time: $O(N)$
Space: $O(K)$