Isolation
Contest: Codeforces Round 542 [Alex Lopashev Thanks-Round] (Div. 1)
Difficulty rating: 2900
Tags: data structures, dp

Find the number of ways to divide an array $$$a$$$ of $$$n$$$ integers into any number of disjoint non-empty segments so that, in each segment, there exist at most $$$k$$$ distinct integers that appear exactly once.

Since the answer can be large, find it modulo $$$998\,244\,353$$$.

Input format:
The first line contains two space-separated integers $$$n$$$ and $$$k$$$ ($$$1 \leq k \leq n \leq 10^5$$$) — the number of elements in the array $$$a$$$ and the restriction from the statement.

The following line contains $$$n$$$ space-separated integers $$$a_1, a_2, \ldots, a_n$$$ ($$$1 \leq a_i \leq n$$$) — elements of the array $$$a$$$.

Output format:
The first and only line contains the number of ways to divide an array $$$a$$$ modulo $$$998\,244\,353$$$.

Constraints:
Time limit: 3.0 s
Memory limit: 256.0 MB

Note:
In the first sample, the three possible divisions are as follows.

- $$$[[1], [1], [2]]$$$
- $$$[[1, 1], [2]]$$$
- $$$[[1, 1, 2]]$$$

Division $$$[[1], [1, 2]]$$$ is not possible because two distinct integers appear exactly once in the second segment $$$[1, 2]$$$.

Example input:
3 1
1 1 2

Example output:
3

Example input:
5 2
1 1 2 1 3

Example output:
14

Editorial:
Let $$$f(l, r)$$$ be the number of integers that appear exactly once in the segment $$$a[l..r]$$$. We can use the following recurrence to compute the answer: $$$dp(n) = \sum\limits_{0 \leq i < n \wedge f(i+1, n) \leq k} dp(i)$$$, where $$$dp(0) = 1$$$.
A naive $$$\mathcal{O}(n^2)$$$ implementation will definitely be too slow. To compute the said recurrence efficiently, we will do as follows.
Preparation
First, let's design an array $$$b$$$ so that $$$f(l, r)$$$ is the sum of elements in segment $$$b[l..r]$$$ for some fixed $$$r$$$. Ideally, it should be easy (i.e., require only a few operations) to transform this array $$$b$$$ into another array $$$b'$$$ that would work with $$$r' = r+1$$$ instead.
One design is as follows. First, let each entry of $$$b$$$ be $$$0$$$. This array now works imaginarily when $$$r = 0$$$. To make it work for $$$r' = r+1$$$, consider the element $$$a_{r+1}$$$.
• If this value appeared before at least twice, set $$$b_i = 0$$$ where $$$i$$$ is the second-last appearance of $$$a_{r+1}$$$ (not counting the appearance at index $$$r+1$$$).
• If this value appeared before at least once, set $$$b_i = -1$$$ where $$$i$$$ is the last appearance of $$$a_{r+1}$$$.
• Finally, set $$$b_{r+1} = 1$$$.
Now, you can see that the sum in the segment $$$b[l..(r+1)]$$$ correctly represents $$$f(l, r+1)$$$ for any $$$1 \leq l \leq r+1$$$!
The Real Thing
Let us divide the array into $$$k$$$ blocks so that each block contains $$$\frac{n}{k}$$$ elements (assume for simplicity that $$$k$$$ divides $$$n$$$). Each block corresponding to some segment $$$[L..R]$$$ should store (1) sum of elements in $$$b[L..R]$$$ (i.e., $$$f(L, R)$$$) and (2), for each $$$-\frac{n}{k} \leq i \leq \frac{n}{k}$$$, $$$q(i) =$$$ sum of $$$dp(x-1)$$$ where $$$f(x, R)$$$ is less than or equal to $$$i$$$. A modification to $$$b_j$$$ for some index $$$j$$$ will require an $$$\mathcal{O}(\frac{n}{k})$$$ update.
With array $$$b$$$ ready for our $$$r$$$, we are ready to compute $$$dp(r)$$$. Let $$$t$$$ be a temporary variable initially equal to $$$0$$$. For each $$$l \leq r$$$ that belongs to the same block as $$$r$$$, add $$$dp(l-1)$$$ to $$$dp(r)$$$ if $$$f(l, r) \leq k$$$, and also add $$$b_l$$$ to $$$t$$$. This runs in $$$\mathcal{O}(\frac{n}{k})$$$. To account for the left possible endpoints from other blocks, for each block, starting from one directly to the left of the block that contains $$$r$$$ to the leftmost block:
• Suppose this block corresponds to the segment $$$[L..R]$$$.
• Let $$$x = k-t$$$.
• If $$$x < -\frac{n}{k}$$$, do nothing.
• If $$$-\frac{n}{k} \leq x \leq \frac{n}{k}$$$, add $$$q(x)$$$ to $$$dp(r)$$$.
• If $$$\frac{n}{k} < x$$$, add $$$q(\frac{n}{k})$$$ to $$$dp(r)$$$.
• Add $$$f(L, R)$$$ to $$$t$$$.
The step above runs in $$$\mathcal{O}(k)$$$. That is, our algorithm takes $$$\mathcal{O}(k + \frac{n}{k})$$$ time to compute $$$dp(r)$$$ for some $$$r$$$.
The time complexity is $$$\mathcal{O}(nk + \frac{n^2}{k})$$$, since there are $$$n$$$ values of $$$r$$$ that we need to compute $$$dp(r)$$$ for.
If we choose $$$k = \sqrt{n}$$$, our solution would run in $$$\mathcal{O}(n\sqrt{n})$$$.