Codeforces Round #829 (Div. 2) E Wish I Knew How to Sort

Wish I Knew How to Sort

概率dp

设计一个 \(dp[i]\) 表示还需要进行 \(i\) 次有效移动的期望次数

何为有效移动？最后的数组是 \(0\) 在左边，\(1\) 在右边

因此只有把两个在错误位置的交换，才算一次有效移动

因此计算出处于左边（本应该是 \(0\) 的位置）的 \(1\)，然后就知道当前需要多少次有效移动，显然 \(dp[0] = 0\)

有状态转移：

\(dp[i] = \frac{i*i}{n*(n-1)/2}dp[i-1] + (1- \frac{i*i}{n*(n-1)/2})dp[i] + 1\)

移项一下：

\(dp[i] = \frac{n*(n-1)/2}{i*i}dp[i-1]\)

#include <iostream>
#include <vector>
using namespace std;
typedef long long ll;
const ll mod = 998244353;

ll qpow(ll x, ll n)
{
    ll ans = 1;
    while(n)
    {
        if(n & 1) ans = x * ans % mod;
        n >>= 1;
        x = x * x % mod;
    }
    return ans % mod;
}

ll inv(ll x)
{
    return qpow(x, mod - 2);
}

int main()
{
    int t;
    cin >> t;
    while(t--)
    {
        int n;
        cin >> n;
        vector<int>a(n);
        for(int i=0; i<n; i++) cin >> a[i];
        int a0 = 0, a1 = 0;
        for(int i=0; i<n; i++) a0 += a[i] == 0;
        for(int i=0; i<a0; i++) a1 += a[i] == 1;
        ll ans = 0;
        for(int i=1; i<=a1; i++) ans = (ans + inv(i) * inv(i) % mod) % mod;
        ans = ans * (n - 1) % mod * n % mod * inv(2) % mod;
        cout << ans << "\n";
    }
    return 0;
}