Skip to content

992. Subarrays with K Different Integers 👍

  • Time: $O(n)$
  • Space: $O(n)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class Solution {
 public:
  int subarraysWithKDistinct(vector<int>& A, int K) {
    return subarrayWithAtMostKDistinct(A, K) -
           subarrayWithAtMostKDistinct(A, K - 1);
  }

 private:
  int subarrayWithAtMostKDistinct(vector<int>& A, int K) {
    int ans = 0;
    vector<int> count(A.size() + 1);

    for (int l = 0, r = 0; r < A.size(); ++r) {
      if (++count[A[r]] == 1)
        --K;
      while (K == -1)
        if (--count[A[l++]] == 0)
          ++K;
      ans += r - l + 1;  // A[l..r], A[l + 1..r], ..., A[r]
    }

    return ans;
  }
};
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
class Solution {
  public int subarraysWithKDistinct(int[] A, int K) {
    return subarraysWithAtMostKDistinct(A, K) - subarraysWithAtMostKDistinct(A, K - 1);
  }

  private int subarraysWithAtMostKDistinct(int[] A, int K) {
    int ans = 0;
    int[] count = new int[A.length + 1];

    for (int l = 0, r = 0; r < A.length; ++r) {
      if (++count[A[r]] == 1)
        --K;
      while (K == -1)
        if (--count[A[l++]] == 0)
          ++K;
      ans += r - l + 1; // A[l..r], A[l + 1..r], ..., A[r]
    }

    return ans;
  }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
class Solution:
  def subarraysWithKDistinct(self, A: List[int], K: int) -> int:
    def subarraysWithAtMostKDistinct(K: int) -> int:
      ans = 0
      count = Counter()

      l = 0
      for r, a in enumerate(A):
        count[a] += 1
        if count[a] == 1:
          K -= 1
        while K < 0:
          count[A[l]] -= 1
          if count[A[l]] == 0:
            K += 1
          l += 1
        ans += r - l + 1  # A[l..r], A[l + 1..r], ..., A[r]

      return ans

    return subarraysWithAtMostKDistinct(K) - subarraysWithAtMostKDistinct(K - 1)