Codeforces Round #829 (Div. 2) E // 概率dp

题目来源：Codeforces Round #829 (Div. 2) E - Wish I Knew How to Sort

题目链接：Problem - E - Codeforces

题意

给定大小为 \(n\) 的仅包含 \(0\)、\(1\) 的数组 \(a\)，每次操作可以随机地选择两个下标 \(i,j\)（\(1 \le i < j \le n\)），若有 \(a_i > a_j\)，则交换 \(a_i\) 和 \(a_j\) 。

问：将整个数组从小到大排序所需要的操作次数的期望。

数据范围：\(1 \le n \le 2·10^5，1 \le \sum{n} \le 2·10^5\).

思路：概率dp

记数组中 \(0\) 的数量为 \(cnt_0\)，\(1\) 的数量为 \(cnt_1\)，则排序后的数组应该是 \(cnt_0\) 个 \(0\)，再有 \(cnt_1\) 个 \(1\)。

因此，在原数组中，位置在 \([1,cnt_0]\) 的 \(1\) 是需要交换的，记其数量为 \(count\)，位置在 \([cnt_0 + 1, n]\) 的 \(0\) 也是需要交换的，且数量等于 \(count\)。若一次操作，选中的下标 \(i,j\)，满足 \(1 \le i \le cnt_0\) 且 \(a_i = 1\)，\(cnt_0 + 1 \le j \le n\) 且 \(a_j = 0\)，则这次操作可以让 \(count\) 减小\(1\)。

记 \(f_i\) 为 当前需要交换的 \(0\)、\(1\) 等于 \(i\) 时，将数组从小到大排序还需要的操作次数的期望。此时进行一次成功操作的可能是：先取出一个需要交换的 \(0\) 的下标，再取出一个需要交换的 \(1\) 的下标，或者反过来。于是 \(f_i\) 转移到 \(f_{i-1}\) 的概率 \(p_i\) 为：\(\large \frac{2·i^2}{n·(n-1)}\)，即有 \(f_i = f_{i - 1} + {\large \frac{1}{p_i}} = f_{i - 1} + {\large \frac{n·(n-1)}{2·i^2}}\).

由于\(f_0=0\)，因此可以从 \(0\) 递推到 \(f_{count}\)。

时间复杂度：\(O(nlogp)\).

代码

#include <bits/stdc++.h>
#define endl '\n'
#define LL long long
using namespace std;

const int mod = 998244353;

int qmi(int m, int k, int p)
{
    int res = 1;
    while(k) {
        if(k & 1) res = (LL)res * m % p;
        m = (LL)m * m % p;
        k >>= 1;
    }
    return res;
}

void solve()
{
    int n;
    cin >> n;

    int cnt[2] = { 0 };
    vector<int> a(n + 1);
    for(int i = 1; i <= n; ++ i) {
        cin >> a[i];
        ++ cnt[a[i]];
    }

    int __count = 0;
    for(int i = 1; i <= cnt[0]; ++ i) {
        __count += a[i] == 1;
    }

    vector<int> f(__count + 1);
    for(int i = 1; i <= __count; ++ i) {
        f[i] = f[i - 1];
        f[i] += n * (n - 1LL) % mod * qmi(2LL * i * i % mod, mod - 2, mod) % mod;
        f[i] %= mod;
    }
    cout << f[__count] << endl;
}

int main()
{
    cin.tie(0);
    ios::sync_with_stdio(false);

    int test;
    cin >> test;
    while(test--) solve();

    return 0;
}