CF2030E MEXimize the Score

· · 题解

CF2030E MEXimize the Score

定义数组 b 的价值为将数组 b 中的元素划分成 k 个非空集合 S_1, ...S_k 后 \sum_{i = 1}^k\text{MEX}(S_i) 的最大值，其中 k 为任意正整数。

给定数组 a，求 a 的所有子序列价值之和。n \leq 2\times 10^5, a_i < n。

Sol：计数 DP

对于数组 a 来说，其价值为 C_0 + \min(C_0, C_1) + ...+\min(C_0, C_1, ...,c_{n - 1})，其中 C_i 代表 i 在 a 中出现次数。

考虑每个 \min(c_0, c_1, ..., c_i) 对答案的贡献，即有多少个子序列满足 \min(c_0, c_1, ..., c_i) = x，其中 c_i 代表在子序列中 i 的出现次数。

定义 dp[i][j] 代表只包含 [0, i] 的元素，且 \min(c_0, c_1, ..., c_i) = j 的子序列数量，则有 dp[0][j] = \binom{C_0}{j}。

考虑转移，对于 dp[i][j] 来说：

虽然看似 O(n^2) 的状态数，但是实际上对于第一维 i 来说，第二维 j 不会超过 \min(C_0, C_1, ..., C_i)，又因为 \sum C_i = n，所以状态数实际上是 O(n) 的；

对于 O(n) 的转移，我们通过后缀和优化成 O(1) 即可。

当然 DP 求的是只包含 [0, i] 的元素满足限制的子序列数量，对于 [i + 1, n - 1] 的元素并没有考虑在内，因为这些元素影响是不影响 \min(c_0, ..., c_i) 的，所以满足 \min(c_0, c_1, ..., c_i) = j 的子序列数量最终为 dp[i][j] \times 2^{C_{i + 1} + C_{i + 2} + ... + C_{n - 1}}。

所以答案为 \sum_{i, j} dp[i][j] \times 2^{C_{i + 1} + C_{i + 2} + ... + C_{n - 1}} \times j。

#include <bits/stdc++.h>
using namespace std;
#define all(x) (x).begin(), (x).end()
#define SZ(x) ((int)(x).size())
// #define int long long
#define pb push_back
#define endl '\n'
using ll = long long;
constexpr ll inf = 1ll << 60;
constexpr ll mod = 998244353;
const int N = 2e5 + 100;

/*
dp[i][j] 代表只包含 [0, i] 的数，且 min(c[0], c[1], ..., c[i]) = j 的子序列数量
显然有 dp[0][j] = C(c[0], j)
*/

inline void add(ll &a, ll b) {
    a += b;
    if (a >= mod) a -= mod;
}
inline ll mul(ll a, ll b) {
    a *= b;
    if (a >= mod) a %= mod;
    return a;
}
inline ll sub(ll a, ll b) {
    a -= b;
    if (a < 0) a += mod;
    return a;
}
inline ll ksm(ll a, int b) {
    ll res = 1;
    for (; b; b >>= 1) {
        if (b & 1) res = mul(res, a);
        a = mul(a, a);
    }
    return res;
}

ll fac[N], inv[N], ifac[N];
inline void comb_init(int n) {
    fac[0] = 1, fac[1] = 1;
    inv[0] = 0, inv[1] = 1;
    ifac[0] = 1, ifac[1] = 1;
    for (int i = 2; i <= n; ++i) {
        fac[i] = fac[i - 1] * i % mod;
        inv[i] = (mod - mod / i) * inv[mod % i] % mod;
        ifac[i] = ifac[i - 1] * inv[i] % mod;
    }
}
inline ll binom(int n, int m) {
    if (n < 0 || m < 0 || n < m) return 0;
    return fac[n] * ifac[m] % mod * ifac[n - m] % mod;
}

int n, a[N];
void solve() {
    cin >> n;
    vector<int> c(n + 10);
    for (int i = 1; i <= n; ++i) {
        cin >> a[i];
        c[a[i]] += 1;
    }
    vector<ll> suf(n + 10);
    for (int i = n - 1; i >= 0; --i) suf[i] = suf[i + 1] + c[i];
    for (int i = 0; i < n; ++i) suf[i] = ksm(2, suf[i]);
    vector<vector<ll>> dp(n + 5);
    dp[0].resize(c[0] + 1, 0);
    for (int j = 0; j <= c[0]; ++j) dp[0][j] = binom(c[0], j);
    vector<ll> s(n + 10, 0);
    for (int i = c[0]; i >= 0; --i) s[i] = (s[i + 1] + dp[0][i]) % mod;
    int m = c[0];
    for (int i = 1; i < n; ++i) {
        m = min(m, c[i]);
        dp[i].resize(m + 1, 0);
        vector<ll> binom_sum(c[i] + 5, 0);
        for (int j = c[i]; j >= 0; --j) {
            binom_sum[j] = (binom_sum[j + 1] + binom(c[i], j)) % mod;
        }
        for (int j = 0; j <= m; ++j) {
            add(dp[i][j], mul(s[j + 1], binom(c[i], j)));
            add(dp[i][j], mul(binom_sum[j], dp[i - 1][j]));
        }
        s[m + 1] = 0;
        for (int j = m; j >= 0; --j) s[j] = (s[j + 1] + dp[i][j]) % mod;
    }
    ll ans = 0;
    for (int i = 0; i < n; ++i) {
        int sz = dp[i].size();
        for (int j = 1; j < sz; ++j) {
            add(ans, mul(mul(dp[i][j], j), (suf[i + 1] == 0 ? 1 : suf[i + 1])));
        }
    }
    cout << ans << endl;
}
signed main(void) {
    comb_init(200010);
    ios::sync_with_stdio(false), cin.tie(0), cout.tie(0);
    int Test = 1;
    cin >> Test;
    while (Test--) solve();
    return 0;
}