Distance to Different
Contest: Educational Codeforces Round 167 (Rated for Div. 2)
Difficulty rating: 2300
Tags: combinatorics, dp, math

Consider an array $$$a$$$ of $$$n$$$ integers, where every element is from $$$1$$$ to $$$k$$$, and every integer from $$$1$$$ to $$$k$$$ appears at least once.

Let the array $$$b$$$ be constructed as follows: for the $$$i$$$-th element of $$$a$$$, $$$b_i$$$ is the distance to the closest element in $$$a$$$ which is not equal to $$$a_i$$$. In other words, $$$b_i = \min \limits_{j \in [1, n], a_j \ne a_i} |i - j|$$$.

For example, if $$$a = [1, 1, 2, 3, 3, 3, 3, 1]$$$, then $$$b = [2, 1, 1, 1, 2, 2, 1, 1]$$$.

Calculate the number of different arrays $$$b$$$ you can obtain if you consider all possible arrays $$$a$$$, and print it modulo $$$998244353$$$.

Input format:
The only line of the input contains two integers $$$n$$$ and $$$k$$$ ($$$2 \le n \le 2 \cdot 10^5$$$; $$$2 \le k \le \min(n, 10)$$$).

Output format:
Print one integer — the number of different arrays $$$b$$$ you can obtain, taken modulo $$$998244353$$$.

Constraints:
Time limit: 2.0 s
Memory limit: 512.0 MB

Example input:
2 2

Example output:
1

Example input:
4 3

Example output:
3

Editorial:
Consider a block of equal elements in $$$a$$$. If we split $$$a$$$ into such blocks, we can consider each block separately — for each element, we need only the distance to the closest border of the block (except for the first and the last block, where we consider only one border). It's quite easy to see that if the length of the block is even (let's say it's $$$2x$$$), then the distances for the elements in this block are $$$[1, 2, 3, \dots, x-1, x, x, x-1, \dots, 3, 2, 1]$$$. And if the length of the block is odd (let's say it's $$$2x-1$$$), the distances are $$$[1, 2, 3, \dots, x-1, x, x-1, \dots, 3, 2, 1]$$$. Now we don't need the actual values in $$$a$$$, we only need the information about the blocks of equal elements.
We need at least $$$k$$$ blocks of equal elements in $$$a$$$, since if the number of blocks is less than $$$k$$$, it's impossible for the array $$$a$$$ to have $$$k$$$ different values (and if there are at least $$$k$$$ blocks, it's possible to assign each block a value so that every integer from $$$1$$$ to $$$k$$$ appears). So, it looks like we need to calculate the number of ways to split the array into at least $$$k$$$ blocks.
However, this only works if there is a bijection between the ways to split the array and the resulting arrays $$$b$$$. Unfortunately, some ways to split the array might result in the same distance array: for example, the distance array $$$[1, 1, 1, 1]$$$ can be obtained either by splitting into four blocks of size $$$1$$$, or into three blocks, the middle of which has size $$$2$$$. So, if there is a block of size $$$2$$$, and it is not the first or the last block in the split, it can be replaced with two blocks of size $$$1$$$, and nothing changes (except for the number of blocks, which goes up).
Thankfully, this is the only such situation we need to handle: any block longer than $$$2$$$ can be uniquely determined by the values in the middle of it (if there is a value $$$x$$$ which is greater than both its neighbors, it is the middle of a block of size $$$2x-1$$$; and if there is a pair of values $$$x$$$ which are greater than both the value to the left and the value to the right, it is the middle of a block of size $$$2x$$$).
So, we need to count the number of ways to split the array of size $$$n$$$ into at least $$$k$$$ blocks so that only the first and the last block can have size $$$2$$$. This can be done with the following dynamic programming: let $$$dp_{i,j}$$$ be the number of ways to split the first $$$i$$$ elements into $$$j$$$ blocks.
This looks like it works in $$$O(n^3)$$$, but we can use two improvements to speed this up:
• the values $$$j > k$$$ are treated identically to the values $$$j=k$$$, so we can limit the value of $$$j$$$ to $$$k$$$ and say that $$$dp_{i,j}$$$ is the number of ways to split the first $$$i$$$ elements into at least $$$j$$$ blocks if $$$j=k$$$;
• we can do transitions in $$$O(1)$$$ instead of $$$O(n)$$$ by using partial sums.
Combining these two optimizations, we get a solution in $$$O(nk)$$$.