MEX Sequences

Created by LXC on Mon Aug 7 08:37:33 2023

https://codeforces.com/problemset/problem/1613/D

ranting: 1900

tag: dp, math

problem

给出一个长度为n的数组a，数组的值在0到n之间。

如果一个序列$s_1, s_2, \cdots, s_n$，每个对于每个$i,1\le i\le n$，满足$|MEX(s_1, s_2, \cdots, s_i)-s_i|\le 1$，那么这个序列是MEX-correct。

现在问数组a有多少个MEX-correct的子序列。

$MEX(s_1, s_2, \cdots, s_i)$代表$s_1, s_2, \cdots, s_i$中第一个未出现的非负数。

$1 \le n \le 5 \cdot 10^5$

solution

对于MEX值为x的MEX-correct的序列，只有两种情况。

• 0 ... 0 1 ... 1 2 ... 2 ... x...x
• 0 ... 0 1 ... 1 2 ... 2 ... x-1...x-1 x+1 ... x+1 x-1 ...

我们可以设$f1_{i,j}$为前i个数种MEX-correct且MEX值为j的第一类序列数，$f2_{i,j}$为前i个数种MEX-correct且MEX值为j的第二类序列数。

看上去状态数是$O(n^2)$，但是状态转移涉及的状态只有$O(n)$个。

假设我们已经计算出了第一类序列状态$f1_{i-1,j}$，对于当前$a_i$：

• $a_i = j-1$，MEX值仍然为j，所以$f1_{i,j} += f1_{i-1,j}$
• $a_i = j$，MEX值为j+1，所以$f1_{i,j+1} += f1_{i-1,j}$
• $a_i = j+1$，MEX值仍然为j，但是转为了第二类序列，所以$f2_{i,j} += f1_{i-1,j}$

假设我们已经计算出了第二类序列状态$f2_{i-1,j}$，对于当前$a_i$：

• $a_i = j-1$，MEX值仍然为j，所以$f2_{i,j} += f2_{i-1,j}$
• $a_i = j$，MEX直接变为j+2，不是MEX-correct。
• $a_i = j+1$，MEX值仍然为j，但是转为了第二类序列，所以$f2_{i,j} += f2_{i-1,j}$

状态数是$O(n)$，每个状态的转移$O(1)$，所以总时间复杂度是$O(n)$。

但是以往的经验告诉我们需要一个$n \times n$ 的二维数组来存储状态。或许可以用n个哈希表存储。

但是最好的方法是滚动数组，由于状态$f_i$只依赖$f_i-1$所以可以直接用一个一维数组不断覆盖之前的值。

要注意转移的次序以防计算状态数错误。例如以下代码不能交换顺序。

f1[x + 1] += f1[x + 1];
f1[x + 1] += f1[x];

code

#include <bits/stdc++.h>
// #define SINGLE_INPUT
#define ll long long
#define ull unsigned long long
#define N 500005
#define MOD 998244353
using namespace std;

#define mint Modint<MOD>
template <const int _MOD>
struct Modint {
    int v;
    Modint() { v = 0; }
    Modint(long long o) { v = o % _MOD; }
    int val() { return v; }
    int pow(long long o) {
        int ret = 1, tmp = v;
        while (o) {
            if (o & 1)
                ret = ((long long)ret * tmp) % _MOD;
            o >>= 1;
            tmp = ((long long)tmp * tmp) % _MOD;
        }
        return ret;
    }
    void operator=(long long o) { v = o % _MOD; }
    bool operator==(long long o) const { return v == o; }
    bool operator==(Modint o) const { return v == o.v; }
    bool operator!=(long long o) const { return v != o; }
    bool operator!=(Modint o) const { return v != o.v; }
    bool operator<(long long o) const { return v < o; }
    bool operator<(Modint o) const { return v < o.v; }
    bool operator>(long long o) const { return v > o; }
    bool operator>(Modint o) const { return v > o.v; }
    bool operator<=(long long o) const { return v <= o; }
    bool operator<=(Modint o) const { return v <= o.v; }
    bool operator>=(long long o) const { return v >= o; }
    bool operator>=(Modint o) const { return v >= o.v; }

    Modint operator+(long long o) const { return *this + Modint(o); }
    Modint operator+(Modint o) const { return ((long long)v + o.v) % _MOD; }
    Modint operator*(long long o) const { return *this * Modint(o); }
    Modint operator*(Modint o) const { return (long long)v * o.v % _MOD; }
    Modint operator-(long long o) const { return *this - Modint(o); }
    Modint operator-(Modint o) const {
        return ((long long)v - o.v + _MOD) % _MOD;
    }
    Modint operator/(long long o) const { return *this / Modint(o); }
    Modint operator/(Modint o) const {
        return ((long long)v * o.pow(_MOD - 2)) % _MOD;
    }

    void operator+=(long long o) { *this = *this + o; }
    void operator+=(Modint o) { *this = *this + o; }
    void operator*=(long long o) { *this = *this * o; }
    void operator*=(Modint o) { *this = *this * o; }
    void operator-=(long long o) { *this = *this - o; }
    void operator-=(Modint o) { *this = *this - o; }
    void operator/=(long long o) { *this = *this / o; }
    void operator/=(Modint o) { *this = *this / o; }

    Modint operator^(long long o) { return Modint(pow(o)); }
    Modint operator^(Modint o) { return Modint(pow(o.v)); }

    template <class T>
    friend bool operator==(T o, Modint u) {
        return u == o;
    }
    template <class T>
    friend Modint operator+(T o, Modint u) {
        return u + o;
    }
    template <class T>
    friend Modint operator*(T o, Modint u) {
        return u * o;
    }
    template <class T>
    friend Modint operator-(T o, Modint u) {
        return Modint(o) - u;
    }
    template <class T>
    friend Modint operator/(T o, Modint u) {
        return Modint(o) / u;
    }

    void operator++() { *this = *this + 1; }
    void operator--() { *this = *this - 1; }
    void operator++(int k) { *this = *this + 1; }
    void operator--(int k) { *this = *this - 1; }

    template <const int T>
    friend std::istream& operator>>(std::istream& in, Modint<T>& modint) {
        ll x;
        in >> x;
        modint = Modint<T>(x);
        return in;
    }

    template <const int T>
    friend std::ostream& operator<<(std::ostream& os, const Modint<T>& modint) {
        os << modint.v;
        return os;
    }
};


void sol() {
    int n;
    cin >> n;
    vector<int> a(n);
    for (int& i : a)
        cin >> i;
    vector<mint> f1(n + 1), f2(n + 1);
    f1[0] = 1;
    for (int i = 0; i < n; i++) {
        int x = a[i];
        if (x + 1 <= n) {
            f1[x + 1] += f1[x + 1];
            f1[x + 1] += f1[x];
            f2[x + 1] += f2[x + 1];
        }
        if (x > 0) {
            f2[x - 1] += f2[x - 1];
            f2[x - 1] += f1[x - 1];
        }
    }
    mint sum = 0;
    for (int i = 1; i <= n; i++) {
        sum += f1[i];
    }
    for (int i = 0; i <= n; i++) {
        sum += f2[i];
    }
    // cout << sum << "\n";
    cout << sum << endl;
}

int main() {
    cout << setprecision(15) << fixed;
    ios::sync_with_stdio(false);
    cin.tie(0);
    cout.tie(0);
#ifndef SINGLE_INPUT
    int t;
    cin >> t;
    while (t--) {
        sol();
    }
#else
    sol();
#endif
    return 0;
}