Editorial for K-Distinct Subarrays
Remember to use this editorial only when stuck, and not to copy-paste code from it. Please be respectful to the problem author and editorialist.
Submitting an official solution before solving the problem yourself is a bannable offence.
Submitting an official solution before solving the problem yourself is a bannable offence.
Approach
It is easier to count the complement.
Let be the number of non-empty subarrays containing at most
distinct values. Then the
required answer is
.
To compute , use a two-pointer window. Expand the right end one position at a time and track
frequencies inside the current window. While the window contains too many distinct values, move
the left end forward until it becomes valid again.
When the current valid window is , there are exactly
valid subarrays ending
at
. Add that quantity to the total.
This runs in time.
Solution (Python)
from collections import defaultdict
n, k = map(int, input().split())
values = list(map(int, input().split()))
def count_at_most(limit: int) -> int:
if limit <= 0:
return 0
freq: dict[int, int] = defaultdict(int)
distinct = 0
left = 0
total = 0
for right, value in enumerate(values):
if freq[value] == 0:
distinct += 1
freq[value] += 1
while distinct > limit:
left_value = values[left]
freq[left_value] -= 1
if freq[left_value] == 0:
distinct -= 1
left += 1
total += right - left + 1
return total
all_subarrays = n * (n + 1) // 2
print(all_subarrays - count_at_most(k - 1))
Comments