LeetCode Solutions
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;
}
};
|
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;
}
}
|
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)
|
