题目
两颗\(n\)个点的树T1和T2，有\(y\)种颜色;
现在给每个点染色，要求公共边端点的颜色相同，求：
1.op=0 , T1和T2都确定，求合法染色方案数；
2.op=1 , T1确定，求所有T2的合法染色方案数的和 ;
3.op=2 , 求所有(T1,T2)合法染色方案数的和；
\(mod \ 998244353\)的值；
$1 \le n \le 10^5 , 1 \le y \lt 998244353 , op = { 0,1,2 } $ ;
题解
op=0
map即可
op=1
Part 1
子集反演:
\[ g(S) = \sum_{S \subset T} f(T) \Leftrightarrow f(S) = \sum_{S \subset T} (-1)^{|T|-|S|}f(T) \\ \]
设\(T_1 \cap T_2 = S\)的方案数为\(f(S)\)，\(T_1 \cap T_2 \supset S\)的方案数为\(g(S)\)；套用得：
\[ \begin{align} ans &= \sum_{S} f(S) y^{n-|S|}\\ &= y^n \sum_{S} y^{-|S|} \sum_{S \subset T}(-1)^{|T|-|S|} g(T) \\ &= y^n \sum_{T} g(T) \sum_{S\subset T}y^{-|S|}(-1)^{|T|-|S|} \\ &= y^n \sum_{T} g(T) \sum_{i=0}^{|T|} (^{|T|}_{\ i})(\frac{1}{y})^{i}(-1)^{|T|-i}\\ &= y^n \sum_{T} g(T) (\frac{1}{y}-1)^{|T|} \\ 令z&= \frac{1}{y}-1\\ ans &= y^n \sum_{T} z^{|T|}g(T) \\ \end{align} \]
Part 2
prufer序列求\(g(T)\)：
\(T\)会形成\(m = n-|T|\)个连通块，设大小为\(a_i\)，\(g(T)\)相当于将其重新连成树的方案数：
\[ g(T) = \sum_{\sum d_i = m-2} (m-2)! \Pi_{i=1}^{m}\frac{a_i^{d_i+1}}{d_i!} \]
考虑先构造一个prufer序列，再在最后依次加入\(1-m\)得到一个长度为\(2m-2\)的度数序列；\(g(T)\)的组合意义就是对这些点染色，\(i\)号连通块有\(a_i\)种颜色可以染；
对前m-2个位置有n中不同的颜色，后m个位置每个有\(a_i\)种颜色，即：
\[ \begin{align} g(T) &= n^{m-2}\Pi_{i=1}^{m} a_i \\ \Rightarrow \ ans &= y^n \sum_{T}z^{n-m} \times n^{m-2}\Pi_{i=1}^{m} a_i\\ \ ans &= \frac{y^n z^n}{n^2} \sum_{T}z^{-m}n^{m}\Pi_{i=1}^{m} a_i \end{align} \]
把\(\Pi_{i=1}^{m} a_i\)看成每个连通块选一个点，\(f_{ \{ u,0/1 \} }\)表示是否选了的贡献即可\(dp\)；
op=2
Part 1
- 同理：
\[ \begin{align} ans &= y^n\sum_{T} z^{|T|} g^2(T) \\ &= \frac{y^n z^n}{n^4} \sum_{T}z^{-m}n^{2m}\Pi_{i=1}^{m} a_i^2 \\ \end{align} \]
- 同理：
\(\sum_{T}\)后面那一坨可以看成是\(m\)个联通块组成的图；
连通块的EGF为\(F(x) \ = \ \frac{n^2}{z} i^2 \times i^{i-2} \times \frac{x^i}{i!} \Rightarrow \frac{n^2i^i}{zi!}x^i\)
所以图的EGF为$G(x) = e^{F(x)} \Rightarrow ans = \frac{y^n z^n n!}{n^4} [x^n]G(x) $
(这个策爷的论文里有讲)
#include<bits/stdc++.h>
#define ll long long
#define mk make_pair
#define mod 998244353
using namespace std;
const int N=100010;
int n,Y,Z,op;
char gc(){
static char*p1,*p2,s[1000000];
if(p1==p2)p2=(p1=s)+fread(s,1,1000000,stdin);
return(p1==p2)?EOF:*p1++;
}
int rd(){
int x=0;char c=gc();
while(c<'0'||c>'9')c=gc();
while(c>='0'&&c<='9')x=(x<<1)+(x<<3)+c-'0',c=gc();
return x;
}
int pw(int x,int y){
int re=1;
if(y<0)y+=mod-1;
while(y){
if(y&1)re=(ll)re*x%mod;
y>>=1;x=(ll)x*x%mod;
}
return re;
}
namespace subtask0{
typedef pair<int,int>pii;
map<pii,bool>mp;
void solve(){
for(int i=1;i<n;++i){
int u=rd(),v=rd();
if(u>v)swap(u,v);
mp[mk(u,v)]=1;
}
int cnt=0;
for(int i=1;i<n;++i){
int u=rd(),v=rd();
if(u>v)swap(u,v);
if(mp[mk(u,v)])cnt++;
}
printf("%d\n",pw(Y,n-cnt));
}
}
namespace subtask1{
int f[N][2],o=1,hd[N],pre,ipre;
struct Edge{int v,nt;}E[N<<1];
void adde(int u,int v){
E[o]=(Edge){v,hd[u]};hd[u]=o++;
E[o]=(Edge){u,hd[v]};hd[v]=o++;
}
void dfs(int u,int fa){
f[u][0]=f[u][1]=pre;
for(int i=hd[u];i;i=E[i].nt){
int v=E[i].v;
if(v==fa)continue;
dfs(v,u);
int t0 = f[u][0] , t1 = f[u][1];
f[u][0] = (1ll * t0 * f[v][1] %mod + 1ll * t0 * f[v][0] %mod * ipre %mod) %mod;
f[u][1] = (1ll * t1 * f[v][1] %mod + 1ll * (1ll*t0*f[v][1]+1ll*t1*f[v][0]) %mod * ipre %mod) %mod;
}
}
void solve(){
if(!Z){cout<<pw(n,n-2)<<endl;return;}
for(int i=1;i<n;++i)adde(rd(),rd());
pre = 1ll * n * pw(Z , mod-2) %mod;
//pre = 1;
ipre = pw(pre , mod-2);
dfs(1,0);
int ans = 1ll * f[1][1] * pw(1ll*Y*Z%mod , n) %mod * pw(1ll*n*n%mod , mod-2) %mod;
cout << ans << endl;
}
}
namespace subtask2{
const int N2=N<<2;
int fac[N2],ifac[N2],iv[N2],a[N2],b[N2],L,G=3,rev[N2];
void ntt(int*A,int len,int f){
for(L=0;(1<<L)<len;++L);
for(int i=0;i<len;++i){
rev[i]=(rev[i>>1]>>1)|((i&1)<<(L-1));
if(i<rev[i])swap(A[i],A[rev[i]]);
}
for(int i=1;i<len;i<<=1){
int wn=pw(G , f * (mod-1)/i/2);
for(int j=0;j<len;j+=i<<1){
int w=1;
for(int k=0;k<i;++k,w=1ll*w*wn%mod){
int x=A[j+k],y=1ll*w*A[j+k+i]%mod;
A[j+k]=(x+y)%mod,A[j+k+i]=(x-y+mod)%mod;
}
}
}
if(!~f)for(int i=0;i<len;++i)A[i]=(ll)A[i]*iv[len]%mod;
}
void cls(int*A,int l,int r){for(int i=l;i<r;++i)A[i]=0;}
void cpy(int*A,int*B,int l){for(int i=0;i<l;++i)A[i]=B[i];}
void der(int*A,int l){for(int i=0;i<l-1;++i)A[i]=1ll*A[i+1]*(i+1)%mod;A[l-1]=0;}
void dif(int*A,int l){for(int i=l-1;i;--i)A[i]=1ll*A[i-1]*iv[i]%mod;A[0]=0;}
void inv(int*A,int*B,int l){
if(l==1){B[0]=1;return;}
static int t[N2];
int len=l<<1;
inv(A,B,l>>1);cls(B,l,len);
cpy(t,A,l);cls(t,l,len);
ntt(B,len,1);ntt(t,len,1);
for(int i=0;i<len;++i)B[i] = B[i] * (2 - 1ll * t[i] * B[i] %mod +mod) %mod;
ntt(B,len,-1);cls(B,l,len);
}
void ln(int*A,int*B,int l){
static int t[N2];
int len=l<<1;
inv(A,B,l);
cpy(t,A,l);cls(t,l,len);der(t,l);
ntt(B,len,1);ntt(t,len,1);
for(int i=0;i<len;++i)B[i] = 1ll * B[i] * t[i] %mod;
ntt(B,len,-1);cls(B,l,len);
dif(B,l);
}
void exp(int*A,int*B,int l){
if(l==1){B[0]=1;return;}
static int t[N2];
int len=l<<1;
exp(A,B,l>>1);cls(B,l,len);
ln(B,t,l);
for(int i=0;i<l;++i)t[i] = ( -t[i] + A[i] + mod)%mod;
t[0]++;//!!!
ntt(B,len,1);ntt(t,len,1);
for(int i=0;i<len;++i)B[i] = 1ll * B[i] * t[i] %mod;
ntt(B,len,-1);cls(B,l,len);
}
void solve(){
if(!Z){cout<<pw(n,2*n-4)<<endl;return ;}
int len=1;for(;len<=2*n;len<<=1);
iv[1]=1;for(int i=2;i<=len;++i)iv[i]=1ll*(mod-mod/i)*iv[mod%i]%mod;
for(int i=fac[0]=ifac[0]=1;i<=len;++i){
ifac[i]=(ll)ifac[i-1]*iv[i]%mod;
fac[i]=(ll)fac[i-1]*i%mod;
}
int tmp = 1ll * n * n %mod * pw(Z,mod-2) %mod;
for(int i=1;i<=n;++i)a[i] = 1ll * tmp * pw(i,i) %mod * ifac[i] %mod ;
exp(a, b, len>>1 );
int ans = 1ll * pw(1ll*Y*Z%mod , n) * pw((ll)n*n%mod*n%mod*n%mod , mod-2) %mod * fac[n] %mod * b[n] %mod;
cout << ans << endl;
}
}
int main(){
freopen("tree.in","r",stdin);
freopen("tree.out","w",stdout);
n=rd();Y=rd();Z=pw(Y,mod-2)-1;op=rd();
if(op==0)subtask0::solve();
else if(op==1)subtask1::solve();
else subtask2::solve();
return 0;
}