洛谷 P5206 - [WC2019]数树（集合反演+NTT）

洛谷题面传送门

神仙多项式+组合数学题，不过还是被我自己想出来了（

首先对于两棵树 \(E_1,E_2\) 而言，为它们填上 \(1\sim y\) 使其合法的方案数显然是 \(y\) 的 \(E_1\cap E_2\) 的连通块次方，又显然 \(E_1,E_2\) 的导出子图是一棵森林，因此 \(E_1\cap E_2\) 连通块个数就是 \(n-|E_1\cap E_2|\)，因此我们要求的答案就是 \(\sum\limits_{E_1}\sum\limits_{E_2}y^{n-|E_1\cap E_2|}\)。

考虑对三个 subtask 分别讨论一下，首先是 \(op=0\)​​，这个就非常 trivial 了，直接暴力用个 set 或 map 之类的东西求解两棵树的交集，白送 6pts（对于非 jxd 选手就白送了 28pts）。

namespace sub0{
	map<int,bool> vis[MAXN+5];
	void solve(){
		int sum=0;
		for(int i=1,u,v;i<n;i++) scanf("%d%d",&u,&v),vis[min(u,v)][max(u,v)]=1;
		for(int i=1,u,v;i<n;i++) scanf("%d%d",&u,&v),sum+=vis[min(u,v)][max(u,v)];
		printf("%d\n",qpow(y,n-sum));
	}
}

接下来考虑 \(op=1\)​ 的情况，注意到这里涉及 \(E_1\)​​ 与 \(E_2\)​ 的交集，而这是不太好直接枚举的，因此考虑按照 LOJ #3399 的讨论进行集合反演，即根据公式 \(f(S)=\sum\limits_{T\subseteq S}f(T)\sum\limits_{T\subseteq S'\subseteq S}(-1)^{|S'|-|T|}\)​，枚举 \(S\subseteq E_1\cap E_2\)​，再枚举 \(S\subseteq T\subseteq E_1\cap E_2\)​，这样我们就可以像莫比乌斯反演那样枚举所有满足 \(T\)​ 是 \(E_2\)​ 边集的一个子集的 \(E_2\)​，具体来说，

\[\begin{aligned}
ans&=\sum\limits_{E_2}y^{n-|E_1\cap E_2|}\\
&=y^n\sum\limits_{E_2}\sum\limits_{S\subseteq E_1\cap E_2}(\dfrac{1}{y})^{|S|}\sum\limits_{T,s.t.S\subseteq T\subseteq E_1\cap E_2}(-1)^{|T|-|S|}\\
&=y^n\sum\limits_{T\subseteq E_1}(\sum\limits_{S\subseteq T}(\dfrac{1}{y})^{|S|}(-1)^{|T|-|S|})\sum\limits_{E_2}[T\subseteq E_2]\\
&=y^n\sum\limits_{T\subseteq E_1}(\sum\limits_{x=0}^{|T|}\dbinom{|T|}{x}(\dfrac{1}{y})^{x}(-1)^{|T|-x})g(T)\\
&=y^n\sum\limits_{T\subseteq E_1}(\dfrac{1-y}{y})^{|T|}g(T)
\end{aligned}
\]

其中 \(g(T)=\sum\limits_{E_2}[T\subseteq E_2]\)，也就是有多少棵树 \(E_2\) 满足 \(T\) 是 \(E_2\)​ 的子集。根据我们幼儿园就学过的结论，\(g(T)=n^{k-2}\prod\limits_{i=1}^ka_i\)，其中 \(k\) 为 \(T\) 形成的连通块个数，\(a_1,a_2,\cdots,a_k\) 分别表示 \(T\) 形成的 \(k\)​ 个连通块的大小。推到这里，式子似乎已经无法继续再化下去，因此考虑组合意义，具体来说我们可以将考虑这样这一个过程：我们在树上选出若干条边，每选出一条边权值乘上 \(\dfrac{1-y}{y}\)​，然后在这些边组成的每个连通块中放上一个球，每放上一个球权值乘上 \(n\)，最后所有方案的权值之和乘上 \(\dfrac{y^n}{n^2}\) 就是答案。

考虑树形 DP，\(dp_{i,0/1}\) 表示目前考虑了 \(i\) 的子树，目前 \(i\) 所在的连通块是否放上了一个球（\(0/1\)）的所有放置方案的权值之和，合并两个子树 \(x,y\)，其中 \(x\) 是 \(y\) 的父亲时就分 \(x\) 所在的连通块原来是否放上一个球，以及 \(x,y\) 之间的边是否被选择即可。

时间复杂度 \(\mathcal O(n)\)

namespace sub1{
	int hd[MAXN+5],to[MAXN*2+5],nxt[MAXN*2+5],ec=0,coef,ivn;
	void adde(int u,int v){to[++ec]=v;nxt[ec]=hd[u];hd[u]=ec;}
	int dp[MAXN+5][2],tmp[2];
	void add(int &x,int v){((x+=v)>=MOD)&&(x-=MOD);}
	void dfs(int x,int f){
		dp[x][0]=1;dp[x][1]=n;
		for(int e=hd[x];e;e=nxt[e]){
			int y=to[e];if(y==f) continue;dfs(y,x);
			memset(tmp,0,sizeof(tmp));
			for(int p=0;p<2;p++) for(int q=0;q<2;q++)
				if(!(p&q)) add(tmp[p|q],1ll*dp[x][p]*dp[y][q]%MOD*coef%MOD);
			for(int p=0;p<2;p++) add(tmp[p],1ll*dp[x][p]*dp[y][1]%MOD);
			for(int p=0;p<2;p++) dp[x][p]=tmp[p];
		} //for(int i=0;i<2;i++) printf("%d %d %d\n",x,i,dp[x][i]);
	}
	void solve(){
		for(int i=1,u,v;i<n;i++) scanf("%d%d",&u,&v),adde(u,v),adde(v,u);
		coef=(MOD-1ll*(y-1)*qpow(y,MOD-2)%MOD)%MOD;dfs(1,0);
//		printf("%d %d\n",coef,ivn);
		printf("%d\n",1ll*dp[1][1]*qpow(y,n)%MOD*qpow(n,MOD-3)%MOD);
	}
}

最后是 \(op=2\) 的情况，此时两棵树都没有给定，不过我们还是可以套用集合反演公式枚举上文中的 \(T\)，类似地有

\[ans=y^n\sum\limits_{T}(\dfrac{1-y}{y})^{|T|}g^2(T)
\]

直接枚举 \(T\) 显然是不行的。考虑枚举 \(T\) 中连通块的个数 \(k\)，这样可以得到：

\[\begin{aligned}
ans&=y^n\sum\limits_{k=1}^n\sum\limits_{a_1+a_2+\cdots+a_k=n,a_i\le a_{i+1}}h(a_1,a_2,\cdots,a_k)(\dfrac{1-y}{y})^{n-k}(n^{k-2}\prod\limits_{i=1}^ka_i)^2
\end{aligned}
\]

其中 \(h(a_1,a_2,\cdots,a_k)\) 表示有多少个森林满足其每个连通块的大小分别为 \(a_1,a_2,\cdots,a_k\)。考虑怎么求 \(h(a_1,a_2,\cdots,a_k)\)​，首先我们选出这 \(k\) 连通块的点集，方案数可以写成一个多重组合数的形式，即 \(\dbinom{n}{a_1,a_2,\cdots,a_k}=\dfrac{n!}{a_1!a_2!\cdots a_k!}\)，其次这 \(k\) 个连通块任意排列得到的图都是相同的，方案数除以 \(k!\)，再其次，每个连通块中连边方案为 \(\prod\limits_{i=1}^ka_i^{a_i-2}\)，因此

\[h(a_1,a_2,\cdots,a_k)=\dfrac{n!}{\prod\limits_{i=1}^ka_i!}·\dfrac{1}{k!}·\prod\limits_{i=1}^ka_i^{a_i-2}
\]

带回去：

\[ans=y^n\sum\limits_{k=1}^n\sum\limits_{a_1+a_2+\cdots+a_k=n,a_i\le a_{i+1}}\dfrac{n!}{\prod\limits_{i=1}^ka_i!}·\dfrac{1}{k!}·\prod\limits_{i=1}^ka_i^{a_i-2}(\dfrac{1-y}{y})^{n-k}(n^{k-2}\prod\limits_{i=1}^ka_i)^2
\]

化简一下：

\[ans=y^nn!\dfrac{1}{n^4}\sum\limits_{k=1}^n\sum\limits_{a_1+a_2+\cdots+a_k=n,a_i\le a_{i+1}}\dfrac{1}{\prod\limits_{i=1}^ka_i!}·\dfrac{1}{k!}·(\dfrac{1-y}{y})^{n-k}n^{2k}\prod\limits_{i=1}^ka_i^{a_i}
\]

\[ans=\dfrac{y^nn!}{n^4}·(\dfrac{1-y}{y})^n·\sum\limits_{k=1}^n\sum\limits_{a_1+a_2+\cdots+a_k=n,a_i\le a_{i+1}}\dfrac{1}{k!}·(\dfrac{y}{1-y})^k·n^{2k}·\prod\limits_{i=1}^k\dfrac{a_i^{a_i}}{a_i!}
\]

如果我们设 \(F(x)=\sum\limits_{n\ge 0}\dfrac{n^n}{n!}x^n\)，那么不难发现后面的 \(\sum\) 可以写成

\[\sum\limits_{k=1}^n\dfrac{(\dfrac{y}{1-y})^k·n^{2k}·[x^n]F^k(x)}{k!}
\]

不难发现 \(k=0\) 和 \(k>n\) 时贡献都是 \(0\)，因此如果设 \(G(x)=\dfrac{y}{1-y}·n^2·F(x)\)，那么式子可写作

\[\sum\limits_{k}\dfrac{G^k(x)}{k!}
\]

然后我们发现这就是

\[\exp G
\]

于是 \(ans=\dfrac{y^nn!}{n^4}·(\dfrac{1-y}{y})^n·[x^n]\exp G\)，多项式 \(\exp\) 一下即可，时间复杂度 \(n\log n\)。

namespace sub2{
	const int pr=3;
	const int ipr=332748118;
	const int MAXP=1<<19;
	int rev[MAXP+5],inv[MAXP+5],fac[MAXP+5],ifac[MAXP+5];
	void init_fac(int n){
		for(int i=(fac[0]=ifac[0]=inv[0]=inv[1]=1)+1;i<=MAXP;i++) inv[i]=1ll*inv[MOD%i]*(MOD-MOD/i)%MOD;
		for(int i=1;i<=n;i++) fac[i]=1ll*fac[i-1]*i%MOD,ifac[i]=1ll*ifac[i-1]*inv[i]%MOD;
	}
	void NTT(vector<int> &a,int len,int type){
		int lg=31-__builtin_clz(len);
		for(int i=0;i<len;i++) rev[i]=(rev[i>>1]>>1)|((i&1)<<lg-1);
		for(int i=0;i<len;i++) if(i<rev[i]) swap(a[i],a[rev[i]]);
		for(int i=2;i<=len;i<<=1){
			int W=qpow((type<0)?ipr:pr,(MOD-1)/i);
			for(int j=0;j<len;j+=i){
				for(int k=0,w=1;k<(i>>1);k++,w=1ll*w*W%MOD){
					int X=a[j+k],Y=1ll*w*a[(i>>1)+j+k]%MOD;
					a[j+k]=(X+Y)%MOD;a[(i>>1)+j+k]=(X-Y+MOD)%MOD;
				}
			}
		} if(!~type){
			int ivn=qpow(len,MOD-2);
			for(int i=0;i<len;i++) a[i]=1ll*a[i]*ivn%MOD;
		}
	}
	vector<int> conv(vector<int> a,vector<int> b){
		int LEN=1;while(LEN<a.size()+b.size()) LEN<<=1;
		a.resize(LEN,0);b.resize(LEN,0);NTT(a,LEN,1);NTT(b,LEN,1);
		for(int i=0;i<LEN;i++) a[i]=1ll*a[i]*b[i]%MOD;
		NTT(a,LEN,-1);return a;
	}
	vector<int> conv(vector<int> a,vector<int> b,int mxl){
		int LEN=1;while(LEN<a.size()+b.size()) LEN<<=1;
		a.resize(LEN,0);b.resize(LEN,0);NTT(a,LEN,1);NTT(b,LEN,1);
		for(int i=0;i<LEN;i++) a[i]=1ll*a[i]*b[i]%MOD;
		NTT(a,LEN,-1);while(a.size()>mxl) a.ppb();return a;
	}
	vector<int> getinv(vector<int> a,int len){
		vector<int> b(len);b[0]=qpow(a[0],MOD-2);
		for(int i=2;i<=len;i<<=1){
//			printf("inv %d\n",i);
			vector<int> c(a.begin(),a.begin()+i);
			vector<int> d(b.begin(),b.begin()+i);
			d=conv(d,d);c=conv(c,d);
			for(int j=0;j<i;j++) b[j]=(2ll*b[j]-c[j]+MOD)%MOD;
		} return b;
	}
	vector<int> direv(vector<int> a,int len){
		vector<int> b(len);
		for(int i=1;i<len;i++) b[i-1]=1ll*a[i]*i%MOD;
		return b;
	}
	vector<int> inter(vector<int> a,int len){
		vector<int> b(len);
		for(int i=1;i<len;i++) b[i]=1ll*a[i-1]*inv[i]%MOD;
		return b;
	}
	vector<int> getln(vector<int> a,int len){
//		printf("%d\n",len);
		vector<int> b=getinv(a,len),_b=direv(a,len);
		b=conv(b,_b);b=inter(b,len);return b;
	}
	vector<int> getexp(vector<int> a,int len){
		vector<int> b(len);b[0]=1;
		for(int i=2;i<=len;i<<=1){
//			printf("exp %d\n",i);
			vector<int> c(b.begin(),b.begin()+i);
			vector<int> d(b.begin(),b.begin()+i);
			d=getln(d,i);
			for(int j=0;j<i;j++) d[j]=(a[j]-d[j]+MOD)%MOD;
			d[0]=(d[0]+1)%MOD;d=conv(c,d);
			for(int j=0;j<i;j++) b[j]=d[j];
		} return b;
	}
	void solve(){
		if(y==1) return printf("%d\n",1ll*qpow(n,MOD-3+n)*qpow(n,MOD-3+n)%MOD),void();
		init_fac(MAXP);
		int LEN=1;while(LEN<=n) LEN<<=1;vector<int> vec(LEN);
		int coef=1ll*n*n%MOD*y%MOD*qpow((1-y+MOD)%MOD,MOD-2)%MOD;
		for(int i=1;i<LEN;i++) vec[i]=1ll*coef*qpow(i,i)%MOD*ifac[i]%MOD;
//		for(int i=0;i<LEN;i++) printf("%d\n",vec[i]);
//		puts("-1");
		vec=getexp(vec,LEN);
//		for(int i=0;i<LEN;i++) printf("%d\n",vec[i]);
		printf("%d\n",1ll*qpow(n,MOD-5)%MOD*vec[n]%MOD*qpow((1-y+MOD)%MOD,n)%MOD*fac[n]%MOD);
	}
}