洛谷P4705 玩游戏 [生成函数，NTT]

传送门

---

这是两个月之前写的题，但没写博客。现在回过头来看一下发现又不会了……

还是要写博客加深记忆。

---

思路

显然期望可以算出总数再乘上\((nm)^{-1}\)。

那么有

\[\begin{align*}
ans_t&=\sum_{i=1}^n \sum_{j=1}^m (a_i+b_j)^t\\
&=\sum_{i=1}^n \sum_{j=1}^m \sum_{k=0}^t {t\choose k} a_i^k b_j^{t-k}\\
&=t!\sum_{k=0}^t (\sum_{i=1}^n \frac{a_i^k}{k!}) (\sum_{j=1}^m \frac{b_j^{t-k}}{(t-k)!})\\
\end {align*}
\]

显然是个卷积的形式，但怎么求\(\sum a_i^k\)呢？

这似乎是个套路。

设多项式

\[f_i(x)=\sum_{n} a_i^nx^n=\frac{1}{1-a_ix}
\]

那么所求即为

\[\sum_{i=1}^nf_i(x)=\sum_{i=1}^n \frac{1}{1-a_ix}
\]

考虑一个式子：

\[\begin{align*}
[\sum_{i=1}^n \ln(1-a_ix)]'&=\sum_{i=1}^n [\ln(1-a_ix)]'\\
&=\sum_{i=1}^n \frac{a_i}{1-a_ix}\\
&=\sum_{i=1}^n \sum_{j} a_i^{j+1}x^j\\
&=\frac{1}{x}\sum_{i=1}^n \sum_{j=1}^{\infty} a_i^jx^j\\
&=\frac 1 x \sum_{i=1}^n [f_i(x)-1]
\end{align*}
\]

好像非常好，但这东西怎么求呢？

设

\[g(x)=\prod_{i=1}^n (1-a_ix)
\]

则有

\[[\sum_{i=1}^n \ln(1-a_ix)]'=[\ln(\prod_{i=1}^n (1-a_ix))]'=[\ln g(x)]'
\]

显然\(g(x)\)可以分治NTT求得，然后就可以求得\(\sum f_i\)，然后再NTT一下就可以得到答案了。

---

代码

#include<bits/stdc++.h>
namespace my_std{
    using namespace std;
    #define rep(i,x,y) for (register int i=(x);i<=(y);++i)
    #define drep(i,x,y) for (register int i=(x);i>=(y);--i)
    #define mod 998244353ll
    #define sz 401010
    typedef long long ll;
    template<typename T>
    inline void read(T& t)
    {
        t=0;char f=0,ch=getchar();
        double d=0.1;
        while(ch>'9'||ch<'0') f|=(ch=='-'),ch=getchar();
        while(ch<='9'&&ch>='0') t=t*10+ch-48,ch=getchar();
        if(ch=='.')
        {
            ch=getchar();
            while(ch<='9'&&ch>='0') t+=d*(ch^48),d*=0.1,ch=getchar();
        }
        t=(f?-t:t);
    }
    template<typename T,typename... Args>
    inline void read(T& t,Args&... args){read(t); read(args...);}
    void file()
    {
        #ifndef ONLINE_JUDGE
        freopen("a.txt","r",stdin);
        #endif
    }
//	inline ll mul(ll a,ll b){ll d=(ll)(a*(double)b/mod+0.5);ll ret=a*b-d*mod;if (ret<0) ret+=mod;return ret;}
}
using namespace my_std;

ll ksm(ll x,int y)
{
    ll ret=1;
    for (;y;y>>=1,x=x*x%mod) if (y&1) ret=ret*x%mod;
    return ret;
}
#define inv(x) ksm(x,mod-2)

int r[sz],limit;
void NTT_init(int n)
{
    limit=1;int l=-1;
    while (limit<=n+n) limit<<=1,++l;
    rep(i,0,limit-1) r[i]=(r[i>>1]>>1)|((i&1)<<l);
}
void NTT(ll *a,int type)
{
    rep(i,0,limit-1) if (r[i]<i) swap(a[i],a[r[i]]);
    for (int mid=1;mid<limit;mid<<=1)
    {
        ll Wn=ksm(3,(mod-1)/mid>>1);if (type==-1) Wn=inv(Wn);
        for (int len=mid<<1,j=0;j<limit;j+=len)
        {
            ll w=1;
            for (int k=0;k<mid;k++,w=w*Wn%mod)
            {
                ll x=a[j+k],y=a[j+k+mid]*w;
                a[j+k]=(x+y)%mod;a[j+k+mid]=(mod*mod+x-y)%mod;
            }
        }
    }
    if (type==1) return;
    ll I=inv(limit);
    rep(i,0,limit-1) a[i]=a[i]*I%mod;
}
namespace GetLn
{
    ll g[sz];
    ll inv[sz];
    ll a[sz],b[sz];
    void work_inv(int n)// inv=g^{-1} (mod x^n)
    {
        if (n==1) return (void)(inv[0]=::inv(g[0]));
        int mid=(n+1)>>1;
        work_inv(mid);
        NTT_init(n);
        rep(i,0,mid-1) a[i]=inv[i];rep(i,mid,limit-1) a[i]=0;
        rep(i,0,n-1) b[i]=g[i];rep(i,n,limit-1) b[i]=0;
        NTT(a,1);NTT(b,1);
        rep(i,0,limit-1) a[i]=a[i]*(mod+2-a[i]*b[i]%mod)%mod;
        NTT(a,-1);
        rep(i,0,n-1) inv[i]=a[i];
    }
    void work1(ll *a,int n){rep(i,0,n-2) a[i]=a[i+1]*(i+1)%mod;a[n-1]=0;}
    void work2(ll *a,int n){drep(i,n,1) a[i]=a[i-1]*::inv(i)%mod;a[0]=0;}
    void Ln(ll *F,int n)// F=ln F (mod x^{n+1})
    {
        memset(g,0,sizeof(g));memset(inv,0,sizeof(inv));
        rep(i,0,n) g[i]=F[i];
        work_inv(n+1);
        work1(g,n);
        NTT_init(n);
        NTT(inv,1);NTT(g,1);
        rep(i,0,limit-1) F[i]=inv[i]*g[i]%mod;
        NTT(F,-1);
        work2(F,n);
        rep(i,n,limit-1) F[i]=0;
    }
}

int S[55],top;// tmp available
ll tmp[55][sz];
void solve(int l,int r,ll *ret,int *a)
{
    if (l==r) return (void)(ret[0]=1,ret[1]=a[l]);
    int mid=(l+r)>>1;
    int ls=S[top--];solve(l,mid,tmp[ls],a);
    int rs=S[top--];solve(mid+1,r,tmp[rs],a);
    NTT_init(r-l+1);
    NTT(tmp[ls],1);NTT(tmp[rs],1);
    rep(i,0,limit-1) ret[i]=tmp[ls][i]*tmp[rs][i]%mod;
    NTT(ret,-1);
    S[++top]=ls;S[++top]=rs;
    rep(i,0,limit-1) tmp[ls][i]=tmp[rs][i]=0;
}

int n,m,T;
ll A[sz],B[sz];
int a[sz>>2],b[sz>>2];
ll fac[sz>>2],_fac[sz>>2];
void init(int n)
{
    fac[0]=_fac[0]=1;
    rep(i,1,n) fac[i]=fac[i-1]*i%mod;
    _fac[n]=inv(fac[n]);
    drep(i,n-1,1) _fac[i]=_fac[i+1]*(i+1)%mod;
}

int main()
{
    file();
    read(n,m);
    rep(i,1,n) read(a[i]);
    rep(i,1,m) read(b[i]);
    read(T);
    init(T);

    rep(i,0,50) S[i]=i;top=50;solve(1,n,A,a);
    GetLn::Ln(A,T+1);
    rep(i,0,50) S[i]=i;top=50;solve(1,m,B,b);
    GetLn::Ln(B,T+1);

    rep(i,1,T) A[i]=A[i]*i%mod,A[i]=(i&1)?A[i]:mod-A[i];
    rep(i,1,T) B[i]=B[i]*i%mod,B[i]=(i&1)?B[i]:mod-B[i];
    A[0]=n;B[0]=m;

    rep(i,0,T) A[i]=_fac[i]*A[i]%mod;
    rep(i,0,T) B[i]=_fac[i]*B[i]%mod;

    NTT_init(T);
    NTT(A,1);NTT(B,1);
    rep(i,0,limit-1) A[i]=A[i]*B[i]%mod;
    NTT(A,-1);

    rep(i,0,T) A[i]=A[i]*fac[i]%mod;

    ll I=inv(1ll*n*m%mod);
    rep(i,1,T) printf("%lld\n",A[i]*I%mod);
}