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:

1
2
3
4
5
6
7
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.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
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.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
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.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
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:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
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)$


Comment