[ZJOI2019]开关（生成函数+背包DP）

注：以下p[i]均表示概率

设F(x)为按i次开关后到达终止状态方案数的EGF，显然F(x)=π(e^p[i]x/p+(-1)^s[i]e^-p[i]x/p)/2，然而方案包含一些多次到达合法方案的状态，需将其排除。n次操作回到原状态的方案数的生成函数G(x)=π(e^p[i]x/p+e^-p[i]x/p)/2。实现时只需要记录F(x)=Σa[i]e^ix/P中a[i]（-P<=i<=P）的系数即可(G(x)也一样)，于是暴力复杂度O(nP)。H(x)为答案的生成函数，显然F、G、H所对应的OGF f、g、h满足f(x)=g(x)h(x)。然后就是EGF向OGF的转化：F(x)=Σa[i]e^ix/P→f(x)=Σa[i]/(1-ix/P)，其中-P<=i<=P，于是此时要求h'(1)，然后根据除法求导公式，可以计算出h'(x)，但x=1时函数不收敛。然后推一下式子就发现本题其实是背包DP了（这里打数学公式太累了就省略一些内容了），复杂度O(nP)

#include<bits/stdc++.h>
using namespace std;
const int N=2e5+7,mid=1e5,mod=998244353;
int n,sp,ans,s[N],a[N],f[N],g[N],tmp[N];
int qpow(int a,int b)
{
    int ret=1;
    while(b)
    {
        if(b&1)ret=1ll*ret*a%mod;
        a=1ll*a*a%mod,b>>=1;
    }
    return ret;
}
int main()
{
    scanf("%d",&n);
    for(int i=1;i<=n;i++)scanf("%d",&s[i]);
    for(int i=1;i<=n;i++)scanf("%d",&a[i]);
    f[mid]=g[mid]=1;
    for(int i=1;i<=n;i++)
    {
        sp+=a[i];
        for(int j=-sp;j<=sp;j++)tmp[j+mid]=(g[j-a[i]+mid]+g[j+a[i]+mid])%mod;
        memcpy(g,tmp,sizeof g);
        for(int j=-sp;j<=sp;j++)tmp[j+mid]=(f[j-a[i]+mid]+(s[i]?mod-f[j+a[i]+mid]:f[j+a[i]+mid]))%mod;
        memcpy(f,tmp,sizeof f);
    }
    for(int i=-sp;i<sp;i++)ans=(ans+1ll*(g[i+mid]-f[i+mid]+mod)*qpow(sp-i,mod-2))%mod;
    ans=1ll*ans*sp%mod;
    printf("%d",ans);
}

View Code