Codeforces 1172 C2 Nauuo and Pictures (hard version) —— 概率DP

This way

题意：

给你n幅画，你每次会取一幅画，取到第i幅画的概率是它的重量/所有重量的和，如果你喜欢这幅画，它的重量会+1，否则会-1，问你最后每幅画的重量的期望。

题解：

这不能和easy一样一个一个去考虑了，用整体法求出最后喜欢的画的重量之和的期望和不喜欢的重量的和的期望，然后将这个比例乘到每个数上面。就比如说一开始喜欢的画的重量之和是$s$,做完之后它的期望是$E(s)$，那对于第i个喜欢的画，它的答案就是$E(s)\ast w[i]/s$.
那么我们怎么求E[s]，只需要用dp[i][j]，表示总共拿i次，拿了j次喜欢的画的概率。然后$E={\sum }_{i=0}^{m}dp[m][i]\ast i$即可。
反之亦然

#include<iostream>
using namespace std;
#define ll long long
const ll mod=998244353;
const int N=2e5+5,M=3e3+5;
ll qpow(ll a,ll b){ll ans=1;for(;b;b>>=1,a=a*a%mod)if(b&1)ans=ans*a%mod;return ans;}
int a[N];
ll w[N],dp[M][M];
int main()
{
    int n,m;
    scanf("%d%d",&n,&m);
    ll l=0,n_l=0;
    for(int i=1;i<=n;i++)scanf("%d",&a[i]);
    for(int i=1;i<=n;i++)
    {
        scanf("%lld",&w[i]);
        a[i]&&(l+=w[i]);
        a[i]||(n_l+=w[i]);
    }
    dp[0][0]=1;
    for(ll i=0;i<=m;i++)
    {
        for(ll j=0;j<=i;j++)
        {
            dp[i+1][j]=(dp[i+1][j]+dp[i][j]*(n_l-(i-j))%mod*qpow(l+n_l+j-(i-j),mod-2))%mod;
            dp[i+1][j+1]=(dp[i+1][j+1]+dp[i][j]*(l+j)%mod*qpow(l+n_l+j-(i-j),mod-2))%mod;
        }
    }
    ll f_l=l,f_n_l=n_l;
    for(ll i=0;i<=m;i++)
    {
        int no=m-i;
        f_l=(f_l+dp[m][i]*i)%mod;
        f_n_l=(f_n_l-dp[m][i]*no%mod+mod)%mod;
    }
    for(int i=1;i<=n;i++)
    {
        if(a[i])
            printf("%lld\n",f_l*qpow(l,mod-2)%mod*w[i]%mod);
        else
            printf("%lld\n",f_n_l*qpow(n_l,mod-2)%mod*w[i]%mod);
    }
    return 0;
}