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})，
由：
设(m=|E_1igcap E_2|)，得到：
也就是说，我们要选出第一颗树的边集的一个子集(E)，设(g(E))为包含这个(E)的生成树个数，则答案为：
设(n-|E|=m)，也就是所选定的边集使得原树分成了(m)的联通块，设每块的点数为(a_1ldots a_m)，然后考虑用(prufer)序列计算答案。
上述式子后面的(sum)是考虑个点在序列中出现了多少次，接下来我们考虑序列上每个位置是哪个点。假设出现的点的序列为(P)。
明显我们有个(O(n^2))的(DP)，(f_{i,j})表示(i)所在的联通块大小为(j)的方案数。
考虑(prod a_i)的意义也可以认为是从每个联通块中选一个点。所以我们优化状态：(f_{i,0/1})表示(i)所在的联通块中有没有选过一个点。这样复杂度就是(O(n))的了。
task2
先特判(y=1)。
延续(task1)的思路。
设(m=n-l)
由：
得：
代码：
#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;
}