Loj #2983. 「WC2019」数树

Loj #2983. 「WC2019」数树

题目背景

白兔喜欢树。

白云喜欢数数。

有 (n) 只鼠，白兔用 (n − 1) 根蓝色绳子把它们连成了一棵树，每根蓝色绳子连着两只鼠，白云用 (n − 1) 根红色绳子把它们连成了一棵树，每根红色绳子连接着两只鼠。

白云要给予每只鼠一个数。这个数可以是 ([1, y]) 中的任意一个整数。

白兔给了白云一个要求：对于两只鼠 (p, q)，若存在一条连接这两只鼠的路径同时属于这两棵树，则 (p) 和 (q) 必须被给予相同的整数。存在一条路径同时属于这两棵树指的是：存在一个序列 ((a_1 = p, a_2, cdots , a_m = q))，使得：对于所有 (i in [1, m − 1])，都有 (a_i) 和 (a_{i+1}) 既有一根红色绳子直接相连也有一根蓝色绳子直接相连。

白云想知道，她有多少种给予数的方案呢？

鼠在不停地挣扎，想要摆脱绳子的束缚。白云还没有思考出来，鼠便把红色绳子都咬断了。

白兔有些气恼，但是他还是想要知道答案。他便问白云：对于所有红色绳子的连接方案，答案的总和（即求所有红色绳子连接方案的给予数方案之和）是多少？

鼠在不停地挣扎，想要摆脱绳子的束缚。白云还没有思考出来，鼠便把蓝色绳子也咬断了。

白兔有些气恼，但是他还是想要知道答案。他便问白云：对于所有红色和蓝色绳子的连接方案，答案的总和（即求所有红色和蓝色绳子连接方案的给予数方案之和）是多少？两个方案不同当且仅当存在至少一对鼠，在两种方案中，这两只鼠之间直接连接的绳子不同（两只鼠之间连接绳子的可能性有 4 种：没有绳子直接连接，只有红色绳子直接连接，只有蓝色绳子直接连接，两种颜色的绳子均直接连接）。

白云哭了。

题目描述

本题包含三个问题：

- 问题 0：已知两棵 (n) 个节点的树的形态（两棵树的节点标号均为 (1) 至 (n)），其中第一棵树是红树，第二棵树是蓝树。要给予每个节点一个 ([1, y]) 中的整数，使得对于任意两个节点 (p, q)，如果存在一条路径 ((a_1 = p, a_2, cdots , a_m = q)) 同时属于这两棵树，则 (p, q) 必须被给予相同的数。求给予数的方案数。

- 存在一条路径同时属于这两棵树的定义见「题目背景」。

- 问题 1：已知蓝树，对于红树的所有 (n^{n−2}) 种选择方案，求问题 0 的答案之和。

- 问题 2：对于蓝树的所有 (n^{n−2}) 种选择方案，求问题 1 的答案之和。

提示：(n) 个节点的树一共有 (n^{n−2}) 种。

在不同的测试点中，你将可能需要回答不同的问题。我们将用 ( ext{op}) 来指代你需要回答的问题编号（对应上述 0、 1、 2）。

由于答案可能很大，因此你只需要输出答案对 (998, 244, 353) 取模的结果即可。

输入格式

从文件 tree.in 中读入数据。

第一行三个用空格隔开的整数 (n, y, ext{op})。

如果 ( ext{op} = 0)，则接下来 (2 imes (n − 1)) 行，前 ((n − 1)) 行每描述一条蓝色绳子，接下来 ((n − 1)) 行每行描述一条红色绳子。

如果 ( ext{op} = 1)，则接下来 ((n − 1)) 行，每行描述一条蓝色绳子。

如果 ( ext{op} = 2)，则接下来没有输入。

描述绳子的各行将包含两个用空格隔开的整数，分别表示被这条绳子连接的两只鼠的编号。鼠的编号是从 (1) 开始的。

输出格式

输出到文件 tree.out 中。

输出一个整数，表示答案对 (998, 244, 353) 取模的结果。

数据范围与提示

(nleq 10^5,1leq y<998244353)

参考博客

task0

假设第一颗树的边集为(E_1)，第二颗树的为(E_2)，答案为(y^{n-|E_1igcap E_2|})

task1

先特判(y=1)的情况。

我们设(Y=frac{1}{y})，

[ans=y^n*sum_{E_2}Y^{|E_1igcap E_2|} ]

由：

[x^k=(x-1+1)^k=sum_{i=0}^kinom{k}{i}(x-1)^i ]

设(m=|E_1igcap E_2|)，得到：

[egin{align} ans&=sum_{E_2}sum_{i=0}^{m}inom{m}{i}(Y-1)^i \ &=sum_{E_2}sum_{E_3subseteq |E_1igcap E_2|}(Y-1)^{|E_3|}\ &=sum_{E_3subseteq E_1}(Y-1)^{|E_2|}sum_{E_3subseteq E_2}1 end{align} ]

也就是说，我们要选出第一颗树的边集的一个子集(E)，设(g(E))为包含这个(E)的生成树个数，则答案为：

[sum_{Esubseteq E_1}(Y-1)^{|E_2|}g(E) ]

设(n-|E|=m)，也就是所选定的边集使得原树分成了(m)的联通块，设每块的点数为(a_1ldots a_m)，然后考虑用(prufer)序列计算答案。

[sum_{sum d_i=m-2}inom{m-2}{a_1ldots a_m}*prod_{i=1}^ma_i^{d_i+1}\ =prod_{i=1}^ma_isum_{sum d_i=m-2}inom{m-2}{a_1ldots a_m}prod_{i=1}^ma_i^{d_i} ]

上述式子后面的(sum)是考虑个点在序列中出现了多少次，接下来我们考虑序列上每个位置是哪个点。假设出现的点的序列为(P)。

[=prod_{i=1}^ma_isum_{P}prod_{i=1}^{m-2} a_{p_i}\ =prod_{i=1}^ma_iprod_{i=1}^{m-2}(sum_{i=1}^ma_i)\ =prod_{i=1}^ma_i*n^{m-2} ]

明显我们有个(O(n^2))的(DP)，(f_{i,j})表示(i)所在的联通块大小为(j)的方案数。

考虑(prod a_i)的意义也可以认为是从每个联通块中选一个点。所以我们优化状态：(f_{i,0/1})表示(i)所在的联通块中有没有选过一个点。这样复杂度就是(O(n))的了。

task2

先特判(y=1)。

延续(task1)的思路。

[ans=sum_E (Y-1)^{|E|}g(E)^2\ =sum_{l=0}^{n-1}(Y-1)^lsum_{|E|=l}g(E)^2 ]

设(m=n-l)

[egin{align} sum_{|E|=l}g(E)^2&=frac{1}{m!}sum_{sum_{i=1}^ma_i=n}inom{n}{a_1ldots a_m}(prod_{i=1}^ma_i^{a_i-2})(prod_{i=1}^ma_i*n^{m-2})^2\ &=frac{n!}{m!n^4}sum_{sum_{i=1}^ma_i=n} prod_{i=1}^mfrac{n^2a_i^{a_i}}{a_i!}\ end{align} ]

[egin{align} ans&=sum_{l=0}^{n-1}(Y-1)^{l}*frac{n!}{m!n^4} sum_{sum_{i=1}^ma_i=n} prod_{i=1}^mfrac{n^2a_i^{a_i}}{a_i!}\ &=sum_{l=0}^{n-1}(Y-1)^l*frac{n!}{n^4}frac{1}{m!}[x^n](sum_{i>0}frac{n^2*i^i}{i!})^m\ &=(Y-1)^nfrac{n!}{n^4}[x^n]sum_{m=1}^nfrac{(sum_{i>0}frac{n^2*i^i}{i!(Y-1)})^m}{m!} end{align} ]

由：

[sum_{i=0}^{infty}frac{f(x)^i}{i!}=exp(f(x)) ]

得：

[ans=(Y-1)^nfrac{n!}{n^4}[x^n]exp((sum_{i>0}frac{n^2*i^i}{i!(Y-1)})) ]

代码：

#include<bits/stdc++.h>
#define ll long long
#define N 200005
 
using namespace std;
inline int Get() {int x=0,f=1;char ch=getchar();while(ch<'0'||ch>'9') {if(ch=='-') f=-1;ch=getchar();}while('0'<=ch&&ch<='9') {x=(x<<1)+(x<<3)+ch-'0';ch=getchar();}return x*f;}

const ll mod=998244353;
ll ksm(ll t,ll x) {
	ll ans=1;
	for(;x;x>>=1,t=t*t%mod)
		if(x&1) ans=ans*t%mod;
	return ans;
}

int n,y,op;
namespace task0 {
	map<int,int>e[N];
	void solve() {
		for(int i=1;i<n;i++) {
			int a=Get(),b=Get();
			e[a][b]=e[b][a]=1;
		}
		int cnt=0;
		for(int i=1;i<n;i++) {
			int a=Get(),b=Get();
			if(e[a][b]) cnt++;
		}
		cout<<ksm(y,n-cnt);
	}
}

struct road {
	int to,nxt;
}s[N<<1];
int h[N],cnt;
void add(int i,int j) {s[++cnt]=(road) {j,h[i]};h[i]=cnt;}

namespace task1 {
	ll f[N][2];
	ll invn,invy;
	void dfs(int v,int fr) {
		f[v][0]=f[v][1]=n;
		for(int i=h[v];i;i=s[i].nxt) {
			int to=s[i].to;
			if(to==fr) continue ;
			dfs(to,v);
			ll t0=f[v][0],t1=f[v][1];
			f[v][0]=f[v][1]=0;
			f[v][0]=t0*f[to][1]%mod;
			f[v][1]=t1*f[to][1]%mod;
			(f[v][0]+=t0*f[to][0]%mod*invn%mod*invy)%=mod;
			(f[v][1]+=(t1*f[to][0]+t0*f[to][1])%mod*invn%mod*invy)%=mod;
		}
	}
	void solve() {
		if(y==1) {cout<<ksm(n,n-2);return ;}
		for(int i=1;i<n;i++) {
			int a=Get(),b=Get();
			add(a,b),add(b,a);
		}
		invn=ksm(n,mod-2);
		invy=ksm(y,mod-2)-1;
		dfs(1,0);
		cout<<f[1][1]*ksm(n,2*(mod-2))%mod*ksm(y,n)%mod;
	}
}

void NTT(ll *a,int d,int flag) {
	static int rev[N<<2];
	static ll G=3;
	int n=1<<d;
	for(int i=0;i<n;i++) rev[i]=(rev[i>>1]>>1)|((i&1)<<d-1);
	for(int i=0;i<n;i++) if(i<rev[i]) swap(a[i],a[rev[i]]);
	for(int s=1;s<=d;s++) {
		int len=1<<s,mid=len>>1;
		ll w=flag==1?ksm(G,(mod-1)/len):ksm(G,mod-1-(mod-1)/len);
		for(int i=0;i<n;i+=len) {
			ll t=1;
			for(int j=0;j<mid;j++,t=t*w%mod) {
				ll u=a[i+j],v=a[i+j+mid]*t%mod;
				a[i+j]=(u+v)%mod;
				a[i+j+mid]=(u-v+mod)%mod;
			}
		}
	}
	if(flag==-1) {
		ll inv=ksm(n,mod-2);
		for(int i=0;i<n;i++) a[i]=a[i]*inv%mod;
	}
}

void Inv(ll *inv,int d,ll *a) {
	static ll A[N<<2];
	if(d==0) {
		inv[0]=ksm(a[0],mod-2);
		return ;
	}
	Inv(inv,d-1,a);
	for(int i=0;i<1<<d+1;i++) A[i]=0;
	for(int i=0;i<1<<d;i++) A[i]=a[i];
	NTT(inv,d+1,1),NTT(A,d+1,1);
	for(int i=0;i<1<<d+1;i++) inv[i]=(2*inv[i]-inv[i]*inv[i]%mod*A[i]%mod+mod)%mod;
	NTT(inv,d+1,-1);
	for(int i=1<<d;i<1<<d+1;i++) inv[i]=0;
}

void Der(ll *a,int d) {
	int n=1<<d;
	for(int i=0;i<n-1;i++) a[i]=a[i+1]*(i+1)%mod;
	a[n-1]=0; 
}

void Int(ll *a,int d) {
	int n=1<<d;
	for(int i=n-1;i>0;i--) a[i]=a[i-1]*ksm(i,mod-2)%mod;
	a[0]=0;
}

void Ln(ll *ln,int d,ll *a) {
	static ll inv[N<<2];
	static ll der[N<<2];
	for(int i=0;i<1<<d+1;i++) der[i]=inv[i]=0;
	for(int i=0;i<1<<d;i++) der[i]=a[i];
	Inv(inv,d,a);
	Der(der,d);
	NTT(inv,d+1,1),NTT(der,d+1,1);
	for(int i=0;i<1<<d+1;i++) ln[i]=der[i]*inv[i]%mod;
	NTT(ln,d+1,-1);
	for(int i=1<<d;i<1<<d+1;i++) ln[i]=0;
	Int(ln,d);
}

void Exp(ll *exp,int d,ll *a) {
	static ll A[N<<2],ln[N<<2];
	if(d==0) {
		exp[0]=1;
		return ;
	}
	Exp(exp,d-1,a);
	for(int i=0;i<1<<d+1;i++) A[i]=ln[i]=0;
	for(int i=1<<d;i<1<<d+1;i++) exp[i]=0;
	for(int i=0;i<1<<d;i++) A[i]=a[i];
	Ln(ln,d,exp);
	NTT(A,d+1,1),NTT(ln,d+1,1),NTT(exp,d+1,1);
	for(int i=0;i<1<<d+1;i++) exp[i]=exp[i]*(1-ln[i]+A[i]+mod)%mod;
	NTT(exp,d+1,-1);
	for(int i=1<<d;i<1<<d+1;i++) exp[i]=0;
}

ll A[N<<2],inv[N<<2];
ll ln[N<<2],ex[N<<2];
namespace task2 {
	ll f[N<<2],g[N<<2];
	ll fac[N],ifac[N];
	void solve() {
		if(y==1) {
			cout<<ksm(n,n-2)*ksm(n,n-2)%mod;
			return ;
		}
		fac[0]=1;
		for(int i=1;i<=n;i++) fac[i]=fac[i-1]*i%mod;
		ifac[n]=ksm(fac[n],mod-2);
		for(int i=n-1;i>=0;i--) ifac[i]=ifac[i+1]*(i+1)%mod;
		
		int d=ceil(log2(n+1));
		ll Y=ksm(y,mod-2)-1;
		ll inv=ksm(Y,mod-2);
		for(int i=1;i<=n;i++) {
			f[i]=1ll*n*n%mod*inv%mod*ksm(i,i)%mod*ifac[i]%mod;
		}
		Exp(g,d,f);
		ll ans=fac[n]*ksm(Y,n)%mod*ksm(n,4*(mod-2))%mod*g[n]%mod;
		cout<<ans*ksm(y,n)%mod;
	}
}

int main() {
	n=Get(),y=Get(),op=Get();
	if(op==0) task0::solve();
	else if(op==1) task1::solve();
	else task2::solve();
	return 0;
}