LibreOJ #2320.「清华集训 2017」生成树计数 prufer序列+组合数学+分治FFT

题意

有$n$个点，每个点有点权$a_i$，定义一棵生成树$T$中第$i$个点的度数为$d_i$，那么该生成树的权值

$$val(T)=(\prod_{i=1}^na_i^{d_i}d_i^m)(\sum_{i=1}^nd_i^m)$$
求所有生成树的价值的和，答案模998244353。
$n\le30000,m\le30$

分析

我们把右边的$\sum$展开后可以看成，生成树其中一个点度数的指数是$2m$其余都是$m$。
先考虑暴力，我们设$f_{i,j}$表示只考虑前面那项，满足前$i$个点在prufer中选了$j$个位置的所有生成树的权值和，$g_{i,j}$则表示前$i$个点且恰好有一个点的指数是$2m$的权值和。
转移比较显然，考虑如何优化。
注意到$m$非常小，那么我们可以把$d_i^m$看成是用$m$种颜色去涂$d_i$个格子，每个格子可以涂任意数量的颜色的方案数。
因为在prufer序列中每个点出现的次数是$d_i-1$，所以我们在后面加上$1到n$，这样每个点出现的次数就是$d_i$了。
这样的话我们就可以枚举每个点涂了哪些格子，然后剩下没涂的格子随便填就好了。
$f_{i,j}$和$g_{i,j}$的定义不变，只是改成了前$i$个点填了$j$个格子的方案。
转移式推出来大概长这样：

$$f_{i,j}=\sum_{k=0}^mf_{i-1,j-k}C_{n-2-j+k}^ka_i^{k+1}(S(m,k)k!+S(m,k+1)(k+1)!)$$
$$g_{i,j}=s1+s2$$
$$s1=\sum_{k=0}^mg_{i-1,j-k}C_{n-2-j+k}^ka_i^{k+1}(S(m,k)k!+S(m,k+1)(k+1)!)$$
$$s2=\sum_{k=0}^{2m}f_{i-1,j-k}C_{n-2-j+k}^ka_i^{k+1}(S(2m,k)k!+S(2m,k+1)(k+1)!)$$
化简一下可以变成 $$F_i=F_{i-1}*A_i$$
$$G_i=G_{i-1}*A_i+F_{i-1}*B_i$$
其中上面每个变量都是多项式，那么用分治FFT优化下就好了。

代码

#include<iostream>
#include<cstdio>
#include<cstdlib>
#include<cstring>
#include<algorithm>

typedef long long LL;

const int N=30005;
const int M=65;
const int MOD=998244353;

int n,m,jc[N],ny[N],a[N],s[M][M],d1[N][M],d2[N][M],L,rev[N*M*2],f[20][N*M*2],g[20][N*M*2];

int ksm(int x,int y)
{
    int ans=1;
    while (y)
    {
        if (y&1) ans=(LL)ans*x%MOD;
        x=(LL)x*x%MOD;y>>=1;
    }
    return ans;
}

void NTT(int *a,int f)
{
    for (int i=0;i<L;i++) if (i<rev[i]) std::swap(a[i],a[rev[i]]);
    for (int i=1;i<L;i<<=1)
    {
        int wn=ksm(3,f==1?(MOD-1)/i/2:MOD-1-(MOD-1)/i/2);
        for (int j=0;j<L;j+=(i<<1))
        {
            int w=1;
            for (int k=0;k<i;k++)
            {
                int u=a[j+k],v=(LL)a[j+k+i]*w%MOD;
                a[j+k]=(u+v)%MOD;a[j+k+i]=(u+MOD-v)%MOD;
                w=(LL)w*wn%MOD;
            }
        }
    }
    int ny=ksm(L,MOD-2);
    if (f==-1) for (int i=0;i<L;i++) a[i]=(LL)a[i]*ny%MOD;
}

void pre()
{
    s[1][1]=1;
    for (int i=2;i<=m*2;i++)
        for (int j=1;j<=i;j++)
            s[i][j]=(s[i-1][j-1]+(LL)s[i-1][j]*j%MOD)%MOD;
    for (int i=1;i<=n;i++)
    {
        d1[i][0]=d2[i][0]=1;
        for (int j=1;j<=m;j++) d1[i][j]=(LL)d1[i][j-1]*a[i]%MOD;
        for (int j=1;j<=m*2;j++) d2[i][j]=(LL)d2[i][j-1]*a[i]%MOD;
    }
    for (int i=1;i<=n;i++)
    {
        for (int j=0;j<=m;j++) d1[i][j]=(LL)d1[i][j]*ny[j]%MOD*((LL)s[m][j]*jc[j]%MOD+(LL)s[m][j+1]*jc[j+1]%MOD)%MOD;
        for (int j=0;j<=m*2;j++) d2[i][j]=(LL)d2[i][j]*ny[j]%MOD*((LL)s[m*2][j]*jc[j]%MOD+(LL)s[m*2][j+1]*jc[j+1]%MOD)%MOD;
    }
}

void solve(int l,int r,int d)
{
    if (l==r)
    {
        for (int i=0;i<=m*2;i++) f[d][i]=d1[l][i],g[d][i]=d2[l][i];
        return;
    }
    int mid=(l+r)/2,L1=m*(mid-l+1)+m,L2=m*(r-mid)+m;
    solve(l,mid,d+1);
    for (int i=0;i<=L1;i++) f[d][i]=f[d+1][i],g[d][i]=g[d+1][i];
    solve(mid+1,r,d+1);
    int lg=0;
    for (L=1;L<=L1+L2;L<<=1,lg++);
    for (int i=0;i<L;i++) rev[i]=(rev[i>>1]>>1)|((i&1)<<(lg-1));
    for (int i=L1+1;i<L;i++) f[d][i]=g[d][i]=0;
    for (int i=L2+1;i<L;i++) f[d+1][i]=g[d+1][i]=0;
    NTT(f[d],1);NTT(g[d],1);NTT(f[d+1],1);NTT(g[d+1],1);
    for (int i=0;i<L;i++) g[d][i]=((LL)f[d][i]*g[d+1][i]+(LL)g[d][i]*f[d+1][i])%MOD,f[d][i]=(LL)f[d][i]*f[d+1][i]%MOD;
    NTT(f[d],-1);NTT(g[d],-1);
}

int main()
{
    scanf("%d%d",&n,&m);
    jc[0]=jc[1]=ny[0]=ny[1]=1;
    for (int i=2;i<=n;i++) jc[i]=(LL)jc[i-1]*i%MOD,ny[i]=(LL)(MOD-MOD/i)*ny[MOD%i]%MOD;
    for (int i=2;i<=n;i++) ny[i]=(LL)ny[i-1]*ny[i]%MOD;
    for (int i=1;i<=n;i++) scanf("%d",&a[i]);
    pre();
    solve(1,n,0);
    int sum=0;
    for (int i=1;i<=n;i++) (sum+=a[i])%=MOD;
    int ans=0;
    for (int i=0;i<=n-2;i++)
        (ans+=(LL)g[0][i]*ksm(sum,n-2-i)%MOD*ny[n-2-i]%MOD)%=MOD;
    for (int i=1;i<=n;i++) ans=(LL)ans*a[i]%MOD;
    printf("%d",(LL)ans*jc[n-2]%MOD);
    return 0;
}