洛谷P5206 [WC2019]数树 [容斥，DP，生成函数，NTT]
Orz神仙题，让我长了许多见识。
长式子警告
思路
y=1
由于y=1时会导致后面一些式子未定义，先抓出来。
printf("%lld",opt==0?1:(opt==1?ksm(n,n-2):ksm(n,2*n-4)))即可。
opt=0
这没什么好说的……统计有多少条边重合即可。
opt=1
为了方便，以下令\(bas=y^{-1}\)。
以下所有集合都为一棵树/一个森林的边集。
先从暴力开始推起：
\[ans=\sum_{T2} bas^{|T1\cap T2|-n}=\frac 1 {bas^n}\sum_{T2} bas^{|T1\cap T2|}=\frac 1 {bas^n} X
\]
注意到\(T1\cap T2\)很烦，考虑取并集的性质，用容斥干掉它。
设
\[f(S)=\sum_{T2} [S=T1\cap T2]bas^{|S|}\\
g(S)=\sum_{T2} [S\subseteq T1\ \&\& \ S\subseteq T2]bas^{|S|}\\
C(S)=\sum_T [S\subseteq T]
\]
那么可以得到：
\[\begin{align*}
g(S)&=[S\subseteq T1]C(S)bas^{|S|}\\
f(S)&=\sum_{S\subseteq T} g(T)bas^{|S|-|T|}(-1)^{|S|-|T|}\\
&=\sum_{S\subseteq T} [T\subseteq T1] C(T) bas^{|S|}(-1)^{|S|-|T|}
\end{align*}
\]
然后我们接着推答案：
\[\begin{align*}
X&=\sum_{S} \sum_{S\subseteq T} [T\subseteq T1] C(T) bas^{|S|}(-1)^{|S|-|T|}\\
&=\sum_{T\subseteq T1} C(T)(-1)^{|T|} \sum_{S\subseteq T} bas^{|S|}(-1)^{|S|}\\
&=\sum_{T\subseteq T1} C(T)(-1)^{|T|} \sum_{i=0}^{|T|} bas^i(-1)^i{|T|\choose i}\\
&=\sum_{S\subseteq T1} C(S) (bas-1)^{|S|}
\end{align*}
\]
这个式子似乎非常优美，但由于有\(C(S)\)的存在，它的复杂度仍然是指数级的。
考虑将\(C(S)\)搞掉。这里有一个结论：
\[C(S)=n^{n-|S|-2}\prod_{i=1}^{n-|S|} a_i
\]
其中\(a_i\)表示第\(i\)个连通块的大小。
这个式子可以用\(prufer\)序列或矩阵树定理证明，然而我懒得写了。
接下来又可以推式子了：
\[\begin{align*}
X&=\sum_{S\subseteq T1} n^{n-|S|-2}(bas-1)^{|S|}\prod_{i=1}^{n-|S|} a_i\\
&=\frac{(bas-1)^n}{n^2}\sum_{S\subseteq T1} \prod_{i=1}^{n-|S|} \frac{n}{bas-1}a_i\\
\end{align*}\\
\]
令\(P=bas-1,k=\frac n P\)，则有
\[X=\frac{P^n}{n^2}\sum_{S\subseteq T1} \prod_{i=1}^{n-|S|} k\times a_i
\]
考虑它的组合意义：
定义原树的一种分割方法的权值：每个连通块中选一个点，选这个点有\(k\)种选法，方案数。
那么上式就是所有分割方法的权值之和。
然后就可以DP了。
设\(f_x\)为\(x\)连通块已经选了点的方案数，\(g_x\)为\(x\)连通块还未选点的方案数，那么有
\[f_x=F\times f_v+F\times g_v+G\times f_v\\
g_x=G\times(f_v+g_v)
\]
其中\(F,G\)表示原来的DP值。
边界是\(f_x=k,g_x=1\)，最后要的答案就是\(f_1\times \frac{P^n}{bas^n n^2}\)。
（注意代码中\(bas=y\)而不是\(y^{-1}\)）
opt=2
发现opt=1的容斥可以复制过来，于是得到
\[\begin{align*}
ans&=\frac{1}{bas^n}\sum_{T1} \sum_{S\subseteq T1} C(S) (bas-1)^{|S|}\\
&=\frac 1 {bas^n}\sum_S (bas-1)^{|S|} [C(S)]^2\\
&=\frac 1 {bas^n}\sum_S (bas-1)^{|S|} n^{2n-2|S|-4}\prod_{i=1}^{n-|S|} a_i^2\\
&=\frac{1}{bas^n}(bas-1)^nn^{-4}\sum_{S} \prod_{i=1}^{n-|S|} \frac{n^2}{(bas-1)}a_i^2\\
&=\frac{1}{bas^n}\frac{(bas-1)^n}{n^4}\sum_{\{a_1,\dots,a_k\},\sum a=n} \prod_{i=1}^k \frac{n^2}{(bas-1)}a_i^{a_i}\\
\end{align*}
\]
后面很像一个完全背包的样子，于是设
\[F(x)=\sum_{a=1}^{\infty} \frac 1 {a!}\frac{n^2a^a}{bas-1}x^a
\]
然后多项式\(\exp\)即可。
代码
#include<bits/stdc++.h>
clock_t t=clock();
namespace my_std{
using namespace std;
#define pii pair<int,int>
#define fir first
#define sec second
#define MP make_pair
#define rep(i,x,y) for (int i=(x);i<=(y);i++)
#define drep(i,x,y) for (int i=(x);i>=(y);i--)
#define go(x) for (int i=head[x];i;i=edge[i].nxt)
#define templ template<typename T>
#define sz 101010
#define mod 998244353ll
typedef long long ll;
typedef double db;
mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());
templ inline T rnd(T l,T r) {return uniform_int_distribution<T>(l,r)(rng);}
templ inline bool chkmax(T &x,T y){return x<y?x=y,1:0;}
templ inline bool chkmin(T &x,T y){return x>y?x=y,1:0;}
templ 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...);}
char __sr[1<<21],__z[20];int __C=-1,__zz=0;
inline void Ot(){fwrite(__sr,1,__C+1,stdout),__C=-1;}
inline void print(register int x)
{
if(__C>1<<20)Ot();if(x<0)__sr[++__C]='-',x=-x;
while(__z[++__zz]=x%10+48,x/=10);
while(__sr[++__C]=__z[__zz],--__zz);__sr[++__C]='\n';
}
void file()
{
#ifndef ONLINE_JUDGE
freopen("a.in","r",stdin);
#endif
}
inline void chktime()
{
#ifndef ONLINE_JUDGE
cout<<(clock()-t)/1000.0<<'\n';
#endif
}
#ifdef mod
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;}
ll inv(ll x){return ksm(x,mod-2);}
#else
ll ksm(ll x,int y){ll ret=1;for (;y;y>>=1,x=x*x) if (y&1) ret=ret*x;return ret;}
#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;
int n,bas,opt;
namespace Solve0
{
set<pii>s;
int MAIN()
{
int x,y,cnt=0;
rep(i,1,n-1) read(x,y),s.insert(MP(x,y));
rep(i,1,n-1) read(x,y),cnt+=(s.find(MP(x,y))!=s.end()||s.find(MP(y,x))!=s.end());
printf("%lld",ksm(bas,n-cnt));
return 0;
}
}
namespace Solve1
{
struct hh{int t,nxt;}edge[sz<<1];
int head[sz],ecnt;
void make_edge(int f,int t)
{
edge[++ecnt]=(hh){t,head[f]};
head[f]=ecnt;
edge[++ecnt]=(hh){f,head[t]};
head[t]=ecnt;
}
ll f[sz],g[sz];
ll P,K;
void dfs(int x,int fa)
{
f[x]=K;g[x]=1;
#define v edge[i].t
go(x) if (v!=fa)
{
dfs(v,x);
ll F=f[x],G=g[x];
f[x]=(F*f[v]%mod+G*f[v]%mod+F*g[v]%mod)%mod;
g[x]=G*(f[v]+g[v])%mod;
}
#undef v
}
int MAIN()
{
int x,y;
P=inv(bas)-1;K=1ll*n*inv(P)%mod;
rep(i,1,n-1) read(x,y),make_edge(x,y);
dfs(1,0);
printf("%lld",f[1]*inv(1ll*n*n%mod)%mod*ksm(P,n)%mod*ksm(bas,n)%mod);
return 0;
}
}
namespace Solve2
{
int r[sz<<2],limit;
void NTT_init(int n)
{
int l=-1;limit=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 (i<r[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=w*a[j+k+mid]%mod;
a[j+k]=(x+y)%mod;a[j+k+mid]=(x-y+mod)%mod;
}
}
}
if (type==1) return;
ll I=inv(limit);
rep(i,0,limit-1) a[i]=a[i]*I%mod;
}
ll tmp1[sz<<2],tmp2[sz<<2],tmp3[sz<<2],tmp4[sz<<2];
void PolyInv(ll *a,ll *f,int n)
{
if (n==1) return (void)(f[0]=inv(a[0]));
int mid=(n+1)>>1;
PolyInv(a,f,mid);
NTT_init(n);
rep(i,0,mid-1) tmp1[i]=f[i];rep(i,mid,limit-1) tmp1[i]=0;
rep(i,0,n-1) tmp2[i]=a[i];rep(i,n,limit-1) tmp2[i]=0;
NTT(tmp1,1);NTT(tmp2,1);
rep(i,0,limit-1) tmp1[i]=tmp1[i]*(mod+2-tmp1[i]*tmp2[i]%mod)%mod;
NTT(tmp1,-1);
rep(i,0,n-1) f[i]=tmp1[i];
rep(i,0,limit-1) tmp1[i]=tmp2[i]=0;
}
void Derivative(ll *a,ll *b,int n){rep(i,0,n-2) b[i]=a[i+1]*(i+1)%mod;b[n-1]=0;}
void Integrate(ll *a,ll *b,int n){drep(i,n-1,1) b[i]=a[i-1]*inv(i)%mod;b[0]=0;}
void PolyLn(ll *a,ll *f,int n)
{
NTT_init(n);
PolyInv(a,tmp3,n);Derivative(a,tmp4,n);
NTT(tmp3,1);NTT(tmp4,1);
rep(i,0,limit-1) tmp1[i]=tmp3[i]*tmp4[i]%mod;
NTT(tmp1,-1);
Integrate(tmp1,f,n);
rep(i,0,limit-1) tmp1[i]=tmp3[i]=tmp4[i]=0;
}
void PolyExp(ll *a,ll *f,int n)
{
if (n==1) return (void)(f[0]=1);
int mid=(n+1)>>1;
PolyExp(a,f,mid);
rep(i,mid,n-1) f[i]=0;
PolyLn(f,tmp2,n);
rep(i,0,n-1) tmp1[i]=f[i];
rep(i,0,n-1) tmp2[i]=(a[i]-tmp2[i]+mod)%mod;
++tmp2[0];
NTT_init(n);
NTT(tmp1,1);NTT(tmp2,1);
rep(i,0,limit-1) tmp1[i]=tmp1[i]*tmp2[i]%mod;
NTT(tmp1,-1);
rep(i,0,n-1) f[i]=tmp1[i];
rep(i,0,limit-1) tmp1[i]=tmp2[i]=0;
}
ll fac[sz],_fac[sz];
void init(){fac[0]=_fac[0]=1;rep(i,1,sz-1) _fac[i]=inv(fac[i]=fac[i-1]*i%mod);}
ll F[sz<<2],Ans[sz<<2];
int MAIN()
{
bas=inv(bas);
init();
rep(i,1,n) F[i]=_fac[i]*n%mod*n%mod*ksm(i,i)%mod*inv(bas-1)%mod;
PolyExp(F,Ans,n+1);
ll ans=Ans[n]*fac[n]%mod*ksm(bas-1,n)%mod*inv(ksm(n,4))%mod;
cout<<ans*ksm(inv(bas),n)%mod;
return 0;
}
}
int main()
{
file();
read(n,bas,opt);
if (bas==1) return printf("%lld",opt==0?1:(opt==1?ksm(n,n-2):ksm(n,2*n-4))),0;
if (opt==0) return Solve0::MAIN();
if (opt==1) return Solve1::MAIN();
if (opt==2) return Solve2::MAIN();
return 0;
}
