Educational Codeforces Round 118 - D. MEX Sequences Solution

View the problem here.

Suppose $x_i$ is a MEX-correct sequence. For all $i(1\le i \le n)$, it is easy to see that $\operatorname{MEX}(x_1, x_2, \dots, x_i)$ is non-decreasing. So the constraint that $|x_i -\operatorname{MEX}(x_1, x_2, \dots, x_i)| \le 1$ basically says that $x_i$ should be “roughly” non-decreasing as well. Further analysis shall formalize this idea.

Let us denote $S_i = \{x_j : 1 \le j \le i\}$ for each position $i$. Then the value of $\operatorname{MEX}(x_1, x_2, \dots, x_i)$ and hence the choice of $x_i$ is determined by the elements in $S_{i-1}$.

Let us start by considering values that $x_1$ may take, with $S_0$ being an empty set.

If $x_1 = 0$, then $\operatorname{MEX}(x_1)=1$ which satisfies our constraint. Then $S_1 = \{0\}$ and let us consider possible values of $x_2$:

$x_2=0$. Then we have $\operatorname{MEX}(x_1, x_2) = 1$. Further, $S_2 = \{0\}$ which is unchanged from $S_1$. Similarly, it is always possible to append $x_{i+1}=x_i$ to the sequence so that $S_{i+1}=S_{i}$ and $\operatorname{MEX}(x_1, x_2, \dots, x_{i+1}) = \operatorname{MEX}(x_1, x_2, \dots, x_{i})$.
$x_2=1$. Then we have $\operatorname{MEX}(x_1, x_2) = 2$. Further, $S_2 = \{0, 1\}$. Similarly, we can see for $S_i = \{0, 1, 2, \dots, j\}$, it is always possible to append $x_{i+1}=j+1$ to the sequence to get $\operatorname{MEX}(x_1, x_2, \dots, x_{i+1})=j+2$.
$x_2 = 2$. Then we have $\operatorname{MEX}(x_1, x_2) = 1$. Further, $S_2 = \{0, 2\}$. In this case, we cannot append $x_3 = 1$ because that would lead to $\operatorname{MEX}(x_1, x_2, x_3) - x_3 = 3 - 1 = 2 > 1$. We cannot append $x_3 \ge 3$ either since that way $\operatorname{MEX}(x_1, x_2, x_3)$ remains $1$ and is at least $2$ from $x_3$. From this point on, the only possible values to append to the sequence are $0$ and $2$. Similarly, for $S_i = \{0, 1, 2, \dots, j, j + 2\}$, where a “jump” has taken place at some previous position, the only possible values of $x_{i+1}$ to append to the sequence are $j$ and $j+2$.
No other values of $x_2$ are possible.

If $x_1=1$, then $\operatorname{MEX}(x_1) = 0$. It is obvious that in this case we can only repeatedly append $1$ to the sequence with no other values possible.

No other values of $x_1$ are possible.

As we can see, there are three patterns of a valid MEX-correct sequence:

1
2
3

(pattern 1) 1+
(pattern 2) 0+, 1+, 2+, ...
(pattern 3) 0+, 1+, 2+, j+, (j+2)+, (j+, (j+2)+)*

where + and * are used similarly as in regular expression.

Let us consider how to solve the original problem, namely counting the number of MEX-correct subsequences of a given sequence $a_i$.

For the first pattern, we only need to count the number of $1$’s in $a_i$, say $m$, then the answer is $2^m-1$.

For the second and third patterns, define three functions as follows:

$f_i(x)$: the number of subsequences of the sequence $a_1, a_2, \dots, a_i$ following pattern 2 and ending with $x$.

$g_i(x)$: the number of subsequences of the sequence $a_1, a_2, \dots, a_i$ following pattern 3, with a “jump” from $x-2$ to $x$ and ending with $x$.

$h_i(x)$: the number of subsequences of the sequence $a_1, a_2, \dots, a_i$ following pattern 3, with a “jump” from $x-2$ to $x$ and ending with $x-2$.

We use three arrays to represent these functions, omitting the $i$-dimension.

Initially, we have $f_0(x)=g_0(x)=h_0(x) = 0$ for all $x$. Then we traverse all the terms of $a_i$ whilst updating the values of the three arrays.

Then for $f$, if $a_i = 0$,
$$
f_i(0) = 2(f_{i-1}(0) + 1)-1.
$$
The explicit solution is $2^r-1$ where $r$ is the number of $0$’s in $a_1, a_2, \dots, a_i$.

Otherwise,
$$
\
f_i(a_i)= f_{i-1}(a_i-1)+2f_{i-1}(a_i).
$$
Because we can either append $a_i$ to a sequence that ends in $a_i - 1$ or $a_i$, or directly use a sequence that already ends in $a_i$.

Notice that all other values of $f_i(x)$ remains the same as $f_{i-1}(x)$, which is why the $i$-dimension is omitted in our representation.

For $g$, if $a_i \ge 2$,
$$
g_i(a_i) = f_{i-1}(a_i-2) + 2g_{i-1}(a_i) + h_{i-1}(a_i).
$$
Because we can

append $a_i$ to a sequence following pattern 2 ending with $a_i-2$;
append $a_i$ to a sequence following pattern 3 ending with $a_i$;
directly use a sequence following pattern 3 ending with $a_i$; or
append $a_i$ to a sequence following pattern 3 ending with $a_i - 2$.

Similarly, for $h$,
$$
h_i(a_i+2) = g_{i-1}(a_i+2)+2h_{i-1}(a_i+2).
$$
Then the answer is
$$
2^m - 1+ \sum_{x=0}^n f(x)+g(x)+h(x) .
$$
(Notice that $0 \le a_i \le n$.)

C++ code:

#include<bits/stdc++.h>
using namespace std;
#define rep(i,from,to) for(int i=(int)(from);i<=(int)(to);++i)
#define rev(i,from,to) for(int i=(int)(from);i>=(int)(to);--i)
#define For(i,to) for(int i=0;i<(int)(to);++i)
typedef long long ll; typedef long double ld; typedef double db;
const int N = 512345;
const ll M = 998244353ll;
ll f[N], g[N], h[N];
ll n, a[N];
void solve()  {
    cin >> n;
    ll res = 1;
    rep(i,1,n) {
        cin >> a[i]; 
        if (a[i] == 1) res = (res * 2) % M; 
    }
    rep(i,0,n) { 
        f[i] = g[i] = h[i] = 0; 
    }
    res--;
    rep(i,1,n) {
        int u = a[i];
        if (u >= 2) {
            g[u] = (f[u - 2] + g[u] * 2 % M + h[u]) % M;
        }
        h[u + 2] = (g[u + 2] + h[u + 2] * 2) % M;
        if (u > 0) {
            f[u] = (f[u - 1] + 2 * f[u] % M) % M;
        }
        else f[0] = ((((f[0] + 1) * 2) % M) - 1 + M) % M;
    }
    rep(i,0,n) {
        res += (f[i] + g[i] + h[i]) % M;
        res %= M;
    }
    cout << ((res+M)%M) << endl;
    
}
int main() {
    int T; cin >> T;
    while (T--) {
        solve();
    }
    return 0;
}