codeforces Nauuo and Pictures

题目描述

题解

暴力永远不是题解的暴力

考虑暴力dp： $f_w[i][j][k]$ 表示目前权值是 $w$ 且 $w$ 是喜欢的，喜欢的和为 $j$ ，不喜欢的为 $k$ ， $i$ 轮后的期望值， $g_w[i][j][k]$ 是类似的，只是 $w$ 是不喜欢的

考虑转移： $f_w[i][j][k]=\frac{w}{j+k} \times f_{w+1}[i-1][j+1][k]+\frac{j-w}{j+k} \times f_w[i-1][j+1][k]+\frac{k}{j+k}\times f_w[i-1][j][k-1]$

然后我们可以证明 $f_w[i][j][k]=wf_1[i][j][k]$ ，证明的话就用数学归纳法，就是 $f_w[i-1][j][k]=wf_1[i-1][j][k]$ ，然后把式子化开就能证了

所以我们现在只需要求 $f_1[i][j][k]$ ，同时我们发现 $j,k$ 状态的加减次数和 $m-i$ 相同，所以我们可以设计dp： $f[i][j]$ 表示喜欢的被选了几次，不喜欢的被选了几次，当前权值为 $1$ ，在 $m-j-k$ 轮后的期望值，然后和上面的转移类似，可以预处理逆元将效率达到 $O(n+mlogP+m^2)$

代码

#include <bits/stdc++.h>
using namespace std;
const int N=1e6+5,M=3005,P=998244353;
int n,m,a[N],b[N],w[N],c[2],f[M][M],g[M][M];
int K(int x,int y){
    int z=1;
    for (;y;y>>=1,x=1ll*x*x%P)
        if (y&1) z=1ll*z*x%P;
    return z;
}
int main(){
    cin>>n>>m;
    for (int i=1;i<=n;i++)
        scanf("%d",&a[i]);
    for (int i=1;i<=n;i++)
        scanf("%d",&w[i]),c[a[i]]+=w[i];
    for (int i=max(m-c[0]-c[1],0);i<=m+m;i++)
        b[i]=K(c[0]+c[1]+i-m,P-2);
    for (int i=m;~i;i--){
        f[i][m-i]=g[i][m-i]=1;
        for (int j=min(m-i-1,c[0]);~j;j--)
            f[i][j]=1ll*(1ll*(c[1]+i+1)*f[i+1][j]%P+1ll*(c[0]-j)*f[i][j+1]%P)*b[i-j+m]%P,
            g[i][j]=1ll*(1ll*(c[1]+i)*g[i+1][j]%P+1ll*(c[0]-j-1)*g[i][j+1]%P)*b[i-j+m]%P;
    }
    for (int i=1;i<=n;i++)
        printf("%lld\n",1ll*w[i]*(a[i]?f[0][0]:g[0][0])%P);
    return 0;
}