At360 E - Random Swaps of Balls

E.Random Swaps of Balls

问题陈述

有 \(N - 1\) 个白球和一个黑球。这些 \(N\) 球排列成一排，黑球最初位于最左边的位置。

高桥正好要进行下面的操作 \(K\) 次。

• 在 \(1\) 和 \(N\) 之间均匀随机地选择一个整数，包括两次。假设 \(a\) 和 \(b\) 是所选的整数。如果是 \(a \neq b\) , 从左边开始交换 \(a\) -th 和 \(b\) -th 两个球。

经过 \(K\) 次操作后，让黑球位于左边第 \(x\) 个位置。求 \(x\) 的期望值，取模 \(998244353\) 。

模数 \(998244353\) 的期望值是多少？可以证明所求的期望值总是有理数。此外，在本题的限制条件下，还可以证明如果用不可约分数 \(\frac{P}{Q}\) 表示这个值，那么就是 \(Q \not \equiv 0 \pmod{998244353}\) 。因此，存在一个唯一的整数 \(R\) ，使得 \(R \times Q \equiv P \pmod{998244353}\), $0 \leq R < 998244353 $ 。报告这个 \(R\) .

样例输入 1

2 1

样例输出 1

499122178

经过一次操作后，黑球位于左起第 1 个位置和第 2 个位置的概率都是 \(\displaystyle \frac{1}{2}\) 。因此，期望值为 \(\displaystyle \frac{3}{2}\) 。

解题思路

记每一位是黑球的概率为\(p_i\)，期望可以用\(\sum\limits_{i=1}^{n}i·p_i\)表示，由于在k次转移后黑球在1号位的概率是特殊的，而所有黑球不在一号位的概率是相等的\((p_2=p_3=p_4=...p_n)\),那我们就可以使用一个较为暴力的转移计算最终黑球在1号位置和在其他位置的概率，最终期望计算为\(E=1·p_1+\sum\limits_{i=2}^{n}i·p_i\)
利用等差数列求和的性质化简后\(E=1·p_1+(1-p_1)·(2+n)/2\)
暴力转移如下，
记当前在1号概率作\(p\),当前不在1号概率作\(q\)，
记上一位在1号概率作\(pre_p\),当前不在1号概率作\(pre_q\)，
则有

\[\begin{cases} p_0 = 1 & \quad \text i = 0\\ q_0 = 0 & \quad \text i = 0\\ p_i = \text{pre}_p \cdot \left( \frac{(n-1)(n-1) + 1 \cdot 1}{n \cdot n} \right) + \text{pre}_q \cdot \left( \frac{2}{n \cdot n} \right) & \quad \text i > 0 \\ q_i = \text{pre}_p \cdot \left( \frac{2(n-1) \cdot 1}{n \cdot n} \right) + \text{pre}_q \cdot \left( \frac{n \cdot n - 2}{n \cdot n} \right) & \quad \text i > 0 \\ \end{cases} \]

注意维护取模998244353暴力转移就好了，复杂度O(k)

AC code

//https://atcoder.jp/contests/abc360/tasks/abc360_e
#include<bits/stdc++.h>
using namespace std;
#define endl '\n'
typedef long long ll;
const ll mod=998244353;
ll q_pow(ll a, ll b) {
    ll s = 1;
    while (b) {
        if (b & 1) {
            s = s * a % mod;
        }
        a = a * a % mod;
        b >>= 1;
    }
    return s;
}

ll inv(ll x) {
    return q_pow(x, mod - 2);
}//乘法逆元

int main() {
    cin.tie(0)->ios::sync_with_stdio(false);
    ll n, k; cin >> n >> k;
    
    ll _no = 0, _yes = 1;

    ll yes_to_no = (n % mod - 1) * 2 % mod * inv(n % mod * n % mod) % mod; //移走的概率
    ll yes_to_yes = (n * n % mod - n * 2 % mod + 2 + mod) % mod * inv(n % mod * n % mod) % mod; //保留在原地的概率
    ll no_to_yes = 2 * inv(n % mod * n % mod) % mod; //移走到1的概率
    ll no_to_no = (n * n % mod - 2 + mod) % mod * inv(n % mod * n % mod) % mod; //保留no的概率

    for (int i = 1; i <= k; i++) {
        ll no = (_yes * yes_to_no % mod + _no * no_to_no % mod) % mod;
        ll yes = (_no * no_to_yes % mod + _yes * yes_to_yes % mod) % mod;
        _no = no, _yes = yes;
    }//进行k次概率转移

    ll E = (_no * (2LL + n) % mod * inv(2LL) % mod + _yes * 1LL) % mod; //计算最终期望
    cout << E << endl;
    return 0;
}