MEX Sequences
Contest: Educational Codeforces Round 118 (Rated for Div. 2)
Difficulty rating: 1900
Tags: dp, math

Let's call a sequence of integers $$$x_1, x_2, \dots, x_k$$$ MEX-correct if for all $$$i$$$ ($$$1 \le i \le k$$$) $$$|x_i - \operatorname{MEX}(x_1, x_2, \dots, x_i)| \le 1$$$ holds. Where $$$\operatorname{MEX}(x_1, \dots, x_k)$$$ is the minimum non-negative integer that doesn't belong to the set $$$x_1, \dots, x_k$$$. For example, $$$\operatorname{MEX}(1, 0, 1, 3) = 2$$$ and $$$\operatorname{MEX}(2, 1, 5) = 0$$$.

You are given an array $$$a$$$ consisting of $$$n$$$ non-negative integers. Calculate the number of non-empty MEX-correct subsequences of a given array. The number of subsequences can be very large, so print it modulo $$$998244353$$$.

Note: a subsequence of an array $$$a$$$ is a sequence $$$[a_{i_1}, a_{i_2}, \dots, a_{i_m}]$$$ meeting the constraints $$$1 \le i_1 < i_2 < \dots < i_m \le n$$$. If two different ways to choose the sequence of indices $$$[i_1, i_2, \dots, i_m]$$$ yield the same subsequence, the resulting subsequence should be counted twice (i. e. two subsequences are different if their sequences of indices $$$[i_1, i_2, \dots, i_m]$$$ are not the same).

Input format:
The first line contains a single integer $$$t$$$ ($$$1 \le t \le 10^5$$$) — the number of test cases.

The first line of each test case contains a single integer $$$n$$$ ($$$1 \le n \le 5 \cdot 10^5$$$).

The second line contains $$$n$$$ integers $$$a_1, a_2, \dots, a_n$$$ ($$$0 \le a_i \le n$$$).

The sum of $$$n$$$ over all test cases doesn't exceed $$$5 \cdot 10^5$$$.

Output format:
For each test case, print a single integer — the number of non-empty MEX-correct subsequences of a given array, taken modulo $$$998244353$$$.

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

Note:
In the first example, the valid subsequences are $$$[0]$$$, $$$[1]$$$, $$$[0,1]$$$ and $$$[0,2]$$$.

In the second example, the valid subsequences are $$$[0]$$$ and $$$[1]$$$.

In the third example, any non-empty subsequence is valid.

Example input:
4
3
0 2 1
2
1 0
5
0 0 0 0 0
4
0 1 2 3

Example output:
4
2
31
7

Editorial:
Let's understand what MEX-correct sequences look like. It turns out there are only two types: $$$[\underline{0, \dots, 0}, \underline{1, \dots, 1}, \dots, \underline{x-1, \dots, x-1}, \underline{x, \dots, x}]$$$ and $$$[\underline{0, \dots, 0}, \underline{1, \dots, 1}, \dots, \underline{x-1, \dots, x-1}, \underline{x+1, \dots, x+1}, \underline{x-1, \dots, x-1}, \underline{x+1, \dots}]$$$. For example, the sequences $$$[0, 0, 1, 1, 1, 2, 3, 3]$$$ and the empty sequence are MEX-correct sequences of the first type, and $$$[1, 1, 1, 1]$$$ and $$$[0, 0, 1, 2, 4, 4, 2, 4, 2, 2]$$$ of the second one.
Let's calculate the dynamic programming $$$\mathrm{dp1}_{i, j}$$$ — the number of MEX-correct subsequences of the first type on the prefix of length $$$i$$$ with $$$\operatorname{MEX}$$$ equal to $$$j$$$ and similarly $$$\mathrm{dp2}_{i, j}$$$ — the number of MEX-correct subsequences of the second type on the prefix of length $$$i$$$ with $$$\operatorname{MEX}$$$ equal to $$$j$$$.
Let's look at the transitions in these dps, and show that there are no other MEX-correct sequences at the same time.
Let the current state be $$$\mathrm{dp1}_{i, j}$$$, and we are trying to add an element equal to $$$x$$$:
• if $$$x < j - 1$$$, then such an element cannot be added;
• if $$$x = j - 1$$$, then the value of $$$\operatorname{MEX}$$$ will not change and the sequence is still of the first type, which means we have a transition to $$$\mathrm{dp1}_{i + 1, j}$$$;
• if $$$x = j$$$, then the value of $$$\operatorname{MEX}$$$ will increase by $$$1$$$, but it will still be of the first type, which means we have a transition to $$$\mathrm{dp1}_{i + 1, j + 1}$$$
• if $$$x = j + 1$$$, then the value of $$$\operatorname{MEX}$$$ will not change, but the sequence will become of the second type, which means we have a transition to $$$\mathrm{dp2}_{i + 1, j}$$$;
• if $$$x > j + 1$$$, then such an element cannot be added.
Let the current state be $$$\mathrm{dp2}_{i, j}$$$, and we are trying to add an element equal to $$$x$$$:
• if $$$x < j - 1$$$, then such an element cannot be added;
• if $$$x = j - 1$$$, then the value of $$$\operatorname{MEX}$$$ will not change and the sequence is still of the second type, which means we have a transition to $$$\mathrm{dp2}_{i + 1, j}$$$;
• if $$$x = j$$$, then such an element cannot be added, because $$$\operatorname{MEX}$$$ will increase by $$$2$$$, which means the absolute difference between $$$\operatorname{MEX}$$$ and $$$x$$$ is greater than $$$1$$$;
• if $$$x = j + 1$$$, then the value of $$$\operatorname{MEX}$$$ will not change and the sequence is still of the second type, which means we have a transition to $$$\mathrm{dp2}_{i + 1, j}$$$;
• if $$$x > j + 1$$$, then such an element cannot be added.
Thus, we considered all possible transitions (adding a new element to the already MEX-correct sequences) and made sure that there are only two types.
While the solution itself works in $$$O(n)$$$ time (because each element $$$x$$$ has $$$O(1)$$$ possible transitions in the dps), it uses $$$O(n^2)$$$ memory, which does not allow us to write that solution as is. However, note that $$$\mathrm{dp1}_i$$$ and $$$\mathrm{dp1}_{i + 1}$$$ (similarly for $$$\mathrm{dp2}$$$) differ in only a few positions (in those that the element $$$a_i$$$ allowed us to make), which means we can store only one-dimensional arrays, $$$\mathrm{dp1}_j$$$ and $$$\mathrm{dp2}_j$$$. Thus, the final complexity of the solution is $$$O(n)$$$.