[题解] AT_abc_f Socks 4

AT传送门

洛谷传送门

题目大意

抽屉里有 \(n\) 种颜色的袜子，其中第 \(i\) 种颜色的袜子有 \(a_i\) 只。

最初，高桥把一只颜色为 \(c\) 的袜子放在抽屉外面，与这些袜子分开，然后重复下面的操作，直到满足终止条件为止：

• 从抽屉中统一随机抽出 \(1\) 只袜子。然后，如果抽屉外的两只袜子颜色相同，则终止操作。否则，选择其中一只袜子放回抽屉。他在选择放回抽屉的袜子时，总是尽量减少未来抽袜子的预期次数。

求期望操作次数除以 \(998244353\) 的余数。

当这个值表示为不可约分数 \(\frac{P}{Q} (Q \gt 0)\) 时，我们有 \(Q \not\equiv 0 \pmod{998244353}\) 。

因此，唯一存在一个满足 \(R \times Q \equiv P \pmod{998244353}, 0 \leq R < 998244353\) 的整数 \(R\) 。求模为 \(998244353\) 的期望值就是求 \(R\)。

思路

求期望，十有八九是dp了。

设 \(dp_i\) 表示抽屉外的袜子颜色为 \(i\) 时，结束游戏的预期次数，\(a_i\) 表示第 \(i\) 中袜子的个数，\(tot=\sum{a_i}-1\) 为抽屉里袜子总数。

如果抽到了两只袜子 \(i,j\)，我们要放回那个袜子？

• 如果 \(a_i\lt a_j\)，把 \(i\) 放回去，因为 \(i\) 数量更多，更有概率抽中另一只一样颜色的袜子。单个贡献为 \(\frac{a_j}{tot}\)，一共有 \(dp_i\) 次，总贡献为 \(\frac{a_j}{tot}\times dp_i=\frac{a_j\cdot dp_i}{tot}\)。
• 如果 \(a_i=a_j\)，游戏结束。贡献为 \(0\)
• 如果 \(a_i\gt a_j\)，把 \(j\) 放回去，同理。贡献为 \(\frac{a_j}{tot}\times dp_j=\frac{a_j\cdot dp_j}{tot}\)。

这样子，\(dp_i = 1 + \displaystyle\sum_{j = 1}^{i-1} \frac{a_j\cdot dp_i}{tot} + \sum_{j = i + 1}^{n} \frac{a_j\cdot dp_j}{tot}\)。这里 \(1\) 表示拿一次，答完一次后把两种情况加起来。

移项得：\((1 - \displaystyle\sum_{j = 1}^{i-1} \frac{a_j}{tot}) dp_i = 1 + \sum_{j = i + 1}^{N} \frac{a_j\cdot dp_j}{tot}\)，就是 \(dp_i=\frac{1 + \displaystyle\sum_{j = i + 1}^{N} \frac{a_j\cdot dp_j}{tot}}{1 - \displaystyle\sum_{j = 1}^{i-1} \frac{a_j}{tot}}\)。

但是时间复杂度是 \(\mathcal{O}(n^2)\)。

由于计算的瓶颈在于 \(\displaystyle\sum_{j=1}^{i-1} \frac{a_j}{tot}\) 和 \(\displaystyle\sum_{j = i + 1}^{n} \frac{a_j\cdot dp_j}{tot}\) ，如果这两个值的值分别为 \(x\) 和 \(y\) ，那么在计算 \(dp_{i-1}\) 之前，我们可以从 \(x\) 中减去 \(\frac{a_{i - 1}}{tot}\) ，再在 \(y\) 中加上 \(\frac{a_i\cdot dp_i}{tot}\) 。这样就能动态维护了。

注意：如果按 \(a_i\) 升序时求 \(dp_i\)，则要知道 \(dp_{i+1\sim n}\)，那么就要倒着求。但是如果 \(a_i\) 降序，就只需要正着求即可。

做法

将 \(a_i\) 倒序，每次动态维护 \(x,y\)，因为有除法，所以考虑使用逆元。

代码

#include <bits/stdc++.h>
#define int long long
#define endl '\n'
#define pi pair<int, int>
// #pragma GCC optimize(2)

using namespace std;
const int INF = INT_MAX;
const int mod = 998244353;
const int N = 3e5 + 10;

int qpow(int a, int b) {
    a %= mod;
    int ans = 1;
    while (b) {
        if (b & 1) ans = ans * a % mod;
        a = a * a % mod;
        b >>= 1;
    }
    return ans % mod;
}

int n, c, a[N], dp[N];
// dp[i]表示持有第i种袜子时，需要操作的次数

signed main() {
    cin.tie(nullptr)->sync_with_stdio(false);
    cin >> n >> c;
    int tot = 0, st = 0;
    for (int i = 1; i <= n; i++) {
        cin >> a[i];
        tot += a[i];
        if (i == c) a[i]++, st = a[i]; // c颜色多一个袜子
    }
    sort(a + 1, a + 1 + n, greater<int>());
    int s = 0, sa = 0; // s维护sum{a_i/tot},sa维护sum{a_i*dp_i/tot}
    for (int i = 1; i <= n; i++) {
        if (i == 1)
            dp[i] = tot * qpow(a[1] - 1, mod - 2) % mod;
        else
            dp[i] = (sa + tot) * qpow(s + a[i] - 1, mod - 2) % mod;
        sa = (sa + dp[i] * a[i] % mod) % mod;
        s = (s + a[i]) % mod;
    }
    for (int i = 1; i <= n; i++)
        if (a[i] == st) return cout << dp[i] << endl, 0;
    return 0;
}