Editorial for K-Distinct Subarrays

Use this editorial only when stuck, and do not 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.

Approach

It is easier to count the complement.

Let $f(t)$ be the number of non-empty subarrays containing at most $t$ distinct values. Then the required answer is

$\frac{n(n + 1)}{2} - f(k - 1)$ .

To compute $f(t)$ , 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 $[l, r]$ , there are exactly $r - l + 1$ valid subarrays ending at $r$ . Add that quantity to the total.

This runs in $O(n)$ time.

Solution (Python)

Code 1

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))

Comments0

No comments yet

Be the first to comment.

Editorial for K-Distinct Subarrays

Approach

Solution (Python)

Comments0

New comment