[LibreOJ #2983]【WC2019】数树【计数】【DP】【多项式】
Description
此题含有三个子问题
问题1：
给出n个点的两棵树，记m为只保留同时在两棵树中的边时连通块的个数，求\(y^m\)
问题2：
给出n个点的一棵树，另外一棵树任意生成，求所有方案总的\(y^m\)的和
问题3：
两棵树均任意生成，求所有方案总的\(y^m\)的和
n<=100000,答案对998244353取模
Solution
问题1：
求出同时在两棵树中的边数v，容易得到m=n-v，直接算即可。
问题2：
记\(z=y^{-1}\)即y在模意义下的逆元，我们就是要求\(z^v\)的总和
直接计算比较困难
考虑这样一个转化\(z^v=\sum\limits_{i=0}^{v}(z-1)^i{v\choose i}\)
考虑它的组合意义，实际上就是我们任意枚举一个重合边集E的子集E0，贡献就是\((z-1)^{|E_0|}\)
我们不妨枚举原树边集的子集\(E_0\)，再计算有多少个生成树包含\(E_0\)
假设\(E_0\)构成了m个连通块，它们的点数分别是\(a_1,a_2,...,a_m\)
那么包含\(E_0\)的生成树个数就是\(n^{m-2}\prod\limits_{i=1}^{m} a_i\)，总的贡献就是\((z-1)^{n-m}n^{m-2}\prod\limits_{i=1}^{m} a_i\)
这东西考虑用prufer序来理解，我们将一个连通块看做一个大的点，那么m-2个连通块在prufer序上都有一个“父亲”，但是这个父亲是一个单点，因此是\(n^{m-2}\)，此外后面的大小之积相当于枚举每个连通块prufer序连出去的边是哪一条。
这显然可以用一个n^2的树形DP来完成，考虑优化
考虑上面的式子的组合意义，我们把常数因子提出来（\((z-1)^nn^{-2}\)），相当于每个连通块都有一个固定的贡献\((z-1)^{-1}n\)，此外再在每个连通块中选择一个点，求总的贡献和。
这就可以用树形DP来做了，记\(F[i][0/1]\)表示当前做完了i的子树，i所在的连通块是否已经贡献过。
直接讨论相邻边是否出现即可转移，时间复杂度\(O(N)\)，非常的巧妙。
问题3：
有了问题2的基础，我们来考虑问题3
不妨同样枚举重合边集的一个子集\(E_0\)，m,a的定义同上
那么贡献变成了$$(z-1){n-m}\left(n\prod a_i \right)^2$$
平方乘到指数上，依然是将常数因子提出来，把贡献分配到每个连通块上
此外由于两棵树都是生成的，因此连通块内部的边也不确定，还要乘上连通块内部的生成树个数
那么一个连通块i的贡献就是
\((z-1)^{-1}n^2a_i^2a_i^{a_i-2}=(z-1)^{-1}n^2a_i^{a_i}\)
那么上面的式子相当于连通块的带标号重复拼接
一个连通块的EGF就是$$F(x)=\sum\limits_{i>0}{(z-1){-1}n2i^i\over i!}$$
拼接以后就是\([x^n]e^{F(x)}\)，乘上前面提出来的贡献就是答案了。
Code
#include <bits/stdc++.h>
#define fo(i,a,b) for(int i=a;i<=b;++i)
#define fod(i,a,b) for(int i=a;i>=b;--i)
#define N 100005
#define LL long long
#define mo 998244353
using namespace std;
int n,tp;
LL m;
LL ksm(LL k,LL n)
{
LL s=1;
for(;n;n>>=1,k=k*k%mo) if(n&1) s=s*k%mo;
return s;
}
namespace subtask1
{
map<int,bool> mp[N];
void solve1()
{
fo(i,1,n-1)
{
int x,y;
scanf("%d%d",&x,&y);
mp[x][y]=mp[y][x]=1;
}
int c=0;
fo(i,1,n-1)
{
int x,y;
scanf("%d%d",&x,&y);
if(mp[x][y]) c++;
}
printf("%lld\n",ksm(m,n-c));
}
}
using namespace subtask1;
namespace subtask2
{
LL f[N][2],rm,pm;
int fs[N],m1,nt[2*N],dt[2*N];
void link(int x,int y)
{
nt[++m1]=fs[x];
dt[fs[x]=m1]=y;
}
void dp(int k,int fa)
{
f[k][0]=1,f[k][1]=0;
for(int i=fs[k];i;i=nt[i])
{
int p=dt[i];
if(p!=fa)
{
dp(p,k);
f[k][1]=(f[k][1]*(f[p][0]+f[p][1])%mo+f[k][0]*f[p][1])%mo;
f[k][0]=f[k][0]*(f[p][0]+f[p][1])%mo;
}
}
f[k][1]=(f[k][1]+f[k][0]*pm%mo*(LL)n)%mo;
}
void solve2()
{
if(m==1)
{
printf("%lld\n",ksm(n,n-2));
return;
}
fo(i,1,n-1)
{
int x,y;
scanf("%d%d",&x,&y);
link(x,y),link(y,x);
}
rm=ksm(m,mo-2);
pm=ksm(rm-1,mo-2);
dp(1,0);
printf("%lld\n",f[1][1]*ksm(ksm(n,mo-2),2)%mo*ksm(rm-1,n)%mo*ksm(m,n)%mo);
}
}
using namespace subtask2;
namespace subtask3
{
#define M 262144
LL wg[M+1],wi[M+1],a[M+1],b[M+1],js[M+1],ny[M+1],ns[M+1];
int bit[M+1],cf[19],l2[M+1];
void prp(int num)
{
fo(i,0,num)
{
wi[i]=wg[i*(M/num)];
bit[i]=(bit[i>>1]>>1)|((i&1)<<(l2[num]-1));
}
}
void NTT(LL *a,bool pd,int num)
{
LL v;
fo(i,0,num-1) if(i<bit[i]) swap(a[i],a[bit[i]]);
for(int h=1,m=2,l=num>>1;m<=num;h=m,m<<=1,l>>=1)
{
int c=(!pd)?l:-l;
for(int j=0;j<num;j+=m)
{
LL *x=a+j,*y=a+h+j,*w=(!pd)?wi:wi+num;
fo(i,0,h-1)
{
v=*y * *w%mo;
*y=(*x-v+mo)%mo,*x=(*x+v)%mo;
x++,y++,w+=c;
}
}
}
if(pd) fo(i,0,num-1) a[i]=a[i]*ny[num]%mo;
}
void getinv(int n,LL *a,LL *b)
{
static LL u1[M+1],u2[M+1];
fo(i,0,cf[l2[n]]-1) b[i]=0;
b[0]=ksm(a[0],mo-2);
for(int m=1,t=2,num=4;m<n;m=t,t=num,num<<=1)
{
prp(num);
fo(i,0,num-1) u1[i]=u2[i]=0;
fo(i,0,m-1) u1[i]=b[i];
fo(i,0,t-1) u2[i]=a[i];
NTT(u1,0,num),NTT(u2,0,num);
fo(i,0,num-1) u1[i]=u1[i]*u1[i]%mo*u2[i]%mo;
NTT(u1,1,num);
fo(i,0,t-1) b[i]=((LL)2*b[i]-u1[i]+mo)%mo;
}
}
void getln(int n,LL *a,LL *b)
{
static LL u1[M+1],u2[M+1];
int num=cf[l2[2*n+1]];
getinv(n+1,a,u1);
fo(i,n+1,num-1) u1[i]=0;
prp(num);
fo(i,0,n-1) u2[i]=a[i+1]*(LL)(i+1)%mo;
fo(i,n,num-1) u2[i]=0;
NTT(u1,0,num),NTT(u2,0,num);
fo(i,0,num-1) u1[i]=u1[i]*u2[i]%mo;
NTT(u1,1,num);
b[0]=0;
fo(i,1,n) b[i]=u1[i-1]*ny[i]%mo;
}
void getexp(int n,LL *a,LL *b)
{
static LL u1[M+1],u2[M+1];
b[0]=1;
fo(i,1,cf[l2[n]]) b[i]=0;
for(int m=1,t=2,num=4;m<=n;m=t,t=num,num<<=1)
{
getln(t-1,b,u1);
fo(i,0,t-1) u1[i]=(-u1[i]+a[i]+mo+mo)%mo;
fo(i,t,num-1) u1[i]=0;
u1[0]++;
fo(i,0,m-1) u2[i]=b[i];
fo(i,m,num-1) u2[i]=0;
prp(num);
NTT(u1,0,num),NTT(u2,0,num);
fo(i,0,num-1) u1[i]=u1[i]*u2[i]%mo;
NTT(u1,1,num);
fo(i,0,t-1) b[i]=u1[i];
}
}
void solve3()
{
if(m==1) {printf("%lld\n",ksm(n,2*n-4));return;}
wg[0]=1,wg[1]=ksm(3,(mo-1)/M);
fo(i,2,M) wg[i]=wg[i-1]*wg[1]%mo;
fo(i,0,18) l2[cf[i]=1<<i]=i;
fod(i,M-1,2) if(!l2[i]) l2[i]=l2[i+1];
js[0]=ns[0]=js[1]=ns[1]=ny[1]=1;
fo(i,2,M) js[i]=js[i-1]*(LL)i%mo,ny[i]=(-ny[mo%i]*(LL)(mo/i)%mo+mo)%mo;
fo(i,2,M) ns[i]=ns[i-1]*ny[i]%mo;
LL rm=ksm(m,mo-2),pm=ksm(rm-1,mo-2);
fo(i,1,n) a[i]=pm*(LL)n%mo*(LL)n%mo*ksm(i,i)%mo*ns[i]%mo;
getexp(n,a,b);
printf("%lld\n",b[n]*ksm(ksm(n,mo-2),4)%mo*ksm(rm-1,n)%mo*ksm(m,n)%mo*js[n]%mo);
}
}
using namespace subtask3;
int main()
{
cin>>n>>m>>tp;
if(tp==0) solve1();
else if(tp==1) solve2();
else solve3();
}
