Codeforces Round #564 (Div. 1) C1/C2. Nauuo and Pictures

题目描述

给定 $n$ 个整数 $w_1,w_2,\cdots, w_n$，以及 $n$ 个取值在 $\{0,1\}$ 的整数 $a_0,a_1,\cdots,a_n$

一共进行 $m$ 轮操作，在每一轮，首先会随机选择一个 $[1,n]$ 的整数，其中选到 $i$ 的概率为 $\frac{w_i}{\sum_{j=1}^{n}w_j}$

如果 $a_i=0$，则 $w_i$ 变为 $w_i-1$，否则 $w_i$ 变为 $w_i+1$

即 $w_i$ 变为 $w_i-(-1)^{a_i}$

求结束后 $w_1,w_2,\cdots,w_n$ 的值的期望

题解

设 $A=\sum_{i=1}^{n}w_i \cdot [a_i=1],B=\sum_{i=1}^{n}w_i \cdot [a_i=0]$

先考虑做 $a_i=1$ 的情况，设 $f_{w}[i][j][k]$ 表示对于权值 $w$，还需要进行 $i$ 轮，其中 $a_i=1$ 的 $w$ 之和为 $j$，且 $a_i=0$ 的 $w$ 之和为 $k$，在进行完后的期望值

首先 $f_{w}[0][j][k]=w$，对于 $i \ge 1$，则有：

$$
f_{w}[i][j][k]= \frac{w}{j+k} f_{w+1}[i-1][j+1][k]+ \frac{j-w}{j+k}f_{w}[i-1][j+1][k]+\frac{k}{j+k}f_{w}[i-1][j][k-1]
$$

由于每次 $j,k$ 只会 $\pm 1$，因此只需要记录分别 $\pm 1$ 了多少次即可，同时如果 $+1$ 了 $x$ 次，且 $-1$ 了 $y$ 次，那么还剩下 $m-x-y$ 轮

也就是可以设 $f'_{w}[i][j]=f_{w}[m-i-j][A+i][B-j]$

可以证明，$f_w[i][j][k]=w \cdot f_1[i][j][k]$，证明暂略

也就是现在只需要求 $f_1[i][j][k]$ 就行了，也就是只需要求 $f_1'[i][j]$，就叫它 $f'[i][j]$ 吧，则：

$$
f'[i][j]=
\begin{cases}
1 & \quad (i+j=m) \\
\frac{2+(A+i-1)}{A+i+B-j} f'[i+1][j]+\frac{B-j}{A+i+B-j}f'[i][j+1] & \quad(i+j<m)
\end{cases}
$$

同理，设 $g'[i][j]$ 表示 $a_i=0$ 的情况，则：

$$
g'[i][j]=
\begin{cases}
1 & \quad(i+j=m) \\
\frac{0+(B-j-1)}{A+i+B-j} g'[i][j+1] + \frac{A+i}{A+i+B-j}g'[i+1][j] & \quad(i+j<m)
\end{cases}
$$

代码

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;

const int mod = 998244353, N = 2e5 + 10, M = 3010;

ll pw(ll a, ll b) {
    ll r = 1;
    for( ; b ; b >>= 1, a = a * a % mod) {
        if(b & 1) {
            r = r * a % mod;
        }
    }
    return r;
}

ll getinv(ll n) {
    return pw(n, mod - 2);
}

ll s[2], sum, f[M][M], g[M][M];
int a[N], w[N], n, m;
void upd(ll &x, ll y) {
    x = (x + y) % mod;
}

int main() {
    scanf("%d%d", &n, &m);
    for(int i = 1 ; i <= n ; ++ i) {
        scanf("%d", &a[i]);
    }
    for(int i = 1 ; i <= n ; ++ i) {
        scanf("%d", &w[i]);
        s[a[i]] += w[i];
        sum += w[i];
    }
    ll A = s[1], B = s[0];
    for(int i = m ; i >= 0 ; -- i) {
        f[i][m - i] = g[i][m - i] = 1;
        for(int j = min(m - i - 1ll, s[0]) ; j >= 0 ; -- j) {
            ll inv = getinv(A + B + i - j);
            upd(f[i][j], ((A + i + 1) * f[i + 1][j] % mod + (B - j) * f[i][j + 1] % mod) * inv % mod);
            upd(g[i][j], ((B - j - 1) * g[i][j + 1] % mod + (A + i) * g[i + 1][j] % mod) * inv % mod);
        }
    }
    for(int i = 1 ; i <= n ; ++ i) {
        printf("%lld\n", (w[i] * (a[i] ? f : g)[0][0] % mod + mod) % mod);
    }
}