MEX counting
Contest: Codeforces Round 758 (Div.1 + Div. 2)
Difficulty rating: 3200
Tags: combinatorics, dp, implementation

For an array $$$c$$$ of nonnegative integers, $$$MEX(c)$$$ denotes the smallest nonnegative integer that doesn't appear in it. For example, $$$MEX([0, 1, 3]) = 2$$$, $$$MEX([42]) = 0$$$.

You are given integers $$$n, k$$$, and an array $$$[b_1, b_2, \ldots, b_n]$$$.

Find the number of arrays $$$[a_1, a_2, \ldots, a_n]$$$, for which the following conditions hold:

- $$$0 \le a_i \le n$$$ for each $$$i$$$ for each $$$i$$$ from $$$1$$$ to $$$n$$$.
- $$$|MEX([a_1, a_2, \ldots, a_i]) - b_i| \le k$$$ for each $$$i$$$ from $$$1$$$ to $$$n$$$.

As this number can be very big, output it modulo $$$998\,244\,353$$$.

Input format:
The first line of the input contains two integers $$$n, k$$$ ($$$1 \le n \le 2000$$$, $$$0 \le k \le 50$$$).

The second line of the input contains $$$n$$$ integers $$$b_1, b_2, \ldots, b_n$$$ ($$$-k \le b_i \le n+k$$$) — elements of the array $$$b$$$.

Output format:
Output a single integer — the number of arrays which satisfy the conditions from the statement, modulo $$$998\,244\,353$$$.

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

Example input:
4 0
0 0 0 0

Example output:
256

Example input:
4 1
0 0 0 0

Example output:
431

Editorial:
Let's count $$$dp[pos][mex][big]$$$ — the number of ways to assign first $$$pos$$$ elements in a way that:
• $$$|MEX([a_1, a_2, \ldots, a_i]) - b_i| \le k$$$ for each $$$i$$$ from $$$1$$$ to $$$pos$$$.
• $$$MEX([a_1, a_2, \ldots, a_{pos}]) = mex$$$
• There are exactly $$$big$$$ distinct elements among $$$a_1, a_2, \ldots, a_{pos}$$$, which are bigger then $$$mex$$$. Let's call them large.
• We don't care about the exact values of large elements, we care only about their positions and who of them is equal to who. For example, arrays $$$[3, 3, 4, 0, 1], [4, 4, 5, 0, 1], [5, 5, 3, 0, 1]$$$ would all be equivalent. We care about the exact values of all other elements.
We only need to consider at most $$$2k+1$$$ possible values of $$$mex$$$ for each $$$pos$$$. Now, let's learn how to transition from $$$pos$$$ to $$$pos+1$$$.
Let's say that now we are at state $$$(pos, mex, big)$$$. Where can we transition after assigning $$$a_{pos+1}$$$?
There are two cases.
• $$$MEX$$$ doesn't change.It happens when $$$a_{pos+1} \neq mex$$$. Then, if $$$a_{pos+1}$$$ is among those $$$big$$$ large elements or is less than $$$mex$$$, the number of large elements won't change. Otherwise, the number of large elements increases by $$$1$$$ (but we don't care about its value other than it being bigger than $$$mex$$$).So, we need to add $$$(mex + big)dp[pos][mex][big]$$$ to $$$dp[pos][mex][big]$$$, and $$$dp[pos][mex][big]$$$ to $$$dp[pos][mex][big+1]$$$.
• $$$MEX$$$ becomes $$$mex_1 > mex$$$.It happens only when $$$a_{pos+1} = mex$$$, and all numbers from $$$mex+1$$$ to $$$mex_1-1$$$ appear among those $$$big$$$ large elements. There are $$$\frac{big!}{(big-(mex_1-mex-1))!}$$$ ways to choose their positions there, and all of them will stop being large, so there will be only $$$big - (mex_1 - mex - 1)$$$ elements remaining.So, we have to add $$$\frac{big!}{(big-(mex_1-mex-1))!}\cdot dp[pos][mex][big]$$$ to $$$dp[pos+1][mex_1][big-(mex_1-mex-1)]$$$ for all valid $$$pos, mex, big, mex_1$$$.
This already gives $$$O(n^2k^2)$$$ solution, as we have $$$n$$$ iterations and $$$O(nk^2)$$$ transitions on each of them. However, this TLEs. So, let's try to optimize our transitions. We will process cases when $$$MEX$$$ doesn't change as before, as it's only $$$O(nk)$$$ transitions. So consider only transitions where $$$MEX$$$ changes.
Note that $$$mex + big$$$ increases exactly by one in such transitions. Also note that we have to multiply by $$$big!$$$ when number of large elements previously was $$$big$$$ and divide by $$$(big - (mex_1 - mex - 1))!$$$ when the number of large elements now becomes $$$big - (mex_1 - mex - 1)$$$.
Then let's only consider states with $$$mex + big = $$$ some fixed $$$T$$$ for now, then we have transitions of sort $$$(mex, T-mex) \to (mex_1, T+1-mex_1)$$$ with coefficients $$$\frac{(T-mex)!}{(T+1-mex_1)!}$$$. So, to find what we have to add to $$$dp[pos+1][mex_1][T+1-mex_1]$$$, we have to find sum of $$$dp[pos][mex][T-mex]\cdot (T-mex)!$$$ over all valid $$$mex<mex_1$$$, and to divide this sum by $$$(T+1-mex_1)!$$$. This is easy to do with prefix sums.
So, the final complexity is $$$O(n^2k)$$$.