Educational Codeforces Round 118 - D. MEX Sequences Solution

View the problem here.

Suppose $x_{i}$ is a MEX-correct sequence. For all $i (1 \leq i \leq n)$ , it is easy to see that $MEX (x_{1}, x_{2}, \dots, x_{i})$ is non-decreasing. So the constraint that $| x_{i} - MEX (x_{1}, x_{2}, \dots, x_{i}) | \leq 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 \leq j \leq i}$ for each position $i$ . Then the value of $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 $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 $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 $MEX (x_{1}, x_{2}, \dots, x_{i + 1}) = MEX (x_{1}, x_{2}, \dots, x_{i})$ .
$x_{2} = 1$ . Then we have $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 $MEX (x_{1}, x_{2}, \dots, x_{i + 1}) = j + 2$ .
$x_{2} = 2$ . Then we have $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 $MEX (x_{1}, x_{2}, x_{3}) - x_{3} = 3 - 1 = 2 > 1$ . We cannot append $x_{3} \geq 3$ either since that way $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 $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) + 2 f_{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} \geq 2$ ,
$g_{i} (a_{i}) = f_{i - 1} (a_{i} - 2) + 2 g_{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) + 2 h_{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 \leq a_{i} \leq 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;
}