题面
题意: 本题包含3个Task：
- Task0：给定两棵树的边集S,T，求(bas^{n-|Sigcap T|})
- Task1：给定一棵树的边集S，求(sum_{T}bas^{n-|Sigcap T|})
- Task2：求(sum_{S}sum_{T}bas^{n-|Sigcap T|})
其中n为点数，(n,bas)都是给定的，(nleq 10^5)，答案对(998244353)取模。
Task0
直接模拟即可。
#include<bits/stdc++.h>
using namespace std;
#define N 200007
#define ll long long
const ll mod=998244353;
struct str
{
int x,y;
};
bool operator <(str a,str b)
{
return a.x<b.x||a.x==b.x&&a.y<b.y;
}
set<str> S;
ll p2(ll x){return x*x%mod;}
ll pw(ll x,ll p)
{
return p?p2(pw(x,p/2))*(p&1?x:1)%mod:1;
}
int main()
{
int x,y,n,op;
ll p;
scanf("%d%lld%d",&n,&p,&op);
for(int i=1;i<n;i++)
{
scanf("%d%d",&x,&y);
if(x>y)swap(x,y);
S.insert({x,y});
}
int cnt=0;
for(int i=1;i<n;i++)
{
scanf("%d%d",&x,&y);
if(x>y)swap(x,y);
if(S.find({x,y})!=S.end())cnt++;
}
ll ans=pw(p,n-cnt);
printf("%lld
",ans);
return 0;
}
Task1
要求(ans=sum_{T}bas^{n-|Sigcap T|})
我们设(ans'=sum_{T}bas^{-|Sigcap T|})，那么(ans=bas^n*ans')
然后我们考虑怎么处理(|Sigcap T|)，可以考虑先枚举它，然后再枚举S和T。
那么我们就需要算出一个(F(E))，表示交集恰好为E的(<S,T>)的数量。
但是这东西不好算，那么我们再设一个(G(E))，表示交集至少为E的(<S,T>)的数量。
那么根据容斥原理我们有：
然后我们有
我们设(P=(bas^{-1}-1))，那么(ans'=sum_{Vsubseteq S}G(V)P^{|V|})
然后我们考虑(G(V))怎么算
至少包含(V)，那么可以看成先将(V)中的边全部连上，形成(n-|V|)个连通块，然后把这些连通块连起来的方案数。
我们设将(V)中的边连上后形成的连通块大小分别为(a_1,a_2...,a_{n-|V|})，那么我们有结论：
证明：
我们考虑如果我们将所有连通块连成一棵树，那么对应的边有多少种方案。
一条连接连通块i和连通块j的边的方案数为(a_i*a_j)
那么将整棵树连出来的方案数为(prod_{i=1}^{n-|V|}a_i^{deg_i})，其中(deg_i)为连通块i在这棵树上的度数。
我们惊奇的发现(deg_i)就是这棵树对应的purfer序列中i出现的次数加1
那么我们就枚举所有的purfer序列，设其为P，数字i在P中的出现次数为(times_i)
发现后面这部分它就等于((a_1+a_2...+a_{n-|V|})^{n-|V|-2})，即(n^{n-|V|-2})
于是就有
证毕。
回到我们的问题，有
为了后续操作的方便，我们将(n,P)放到连乘里面
然后我们设(K=nP^{-1},ans''=sum_{Vsubseteq S}prod _{i=1}^{n-|V|}a_iK)，那么(ans'=ans''*n^{-2}P^n)
我们考虑(ans'')怎么算
容易想到用树上dp解决，设(dp[v][i])为在v子树内，与v相连的连通块大小为i的所有方案的贡献和（注意这里的贡献不包含与v相连的那个连通块）
另外我们规定(dp[v][0])即为v子树内的答案。
那么有转移方程
初始状态为(dp[v][1]=1)
最后的答案就是(dp[1][1])
但是这样复杂度是(O(n^2))的，考虑怎么优化
我们其实没有必要将每一个(dp[v][i])都算出来，我们只需要知道它们整体的一个值就可以。
于是我们设(f[v]=sum_{i=0}^{n}dp[v][i])，(g[v]=dp[v][0]=sum_{i=1}^ndp[v][i]*i*K)
那么在转移过程中有
最后还要(f[v]+=g[v])
初始状态为(f[v]=1,g[v]=K)
这样就可以(O(n))转移了
最后的答案为(bas^nn^{-2}P^{n}g[1])
#include<bits/stdc++.h>
using namespace std;
#define N 200007
#define M 400007
#define ll long long
const ll mod=998244353;
int f[N],g[N],n,P,K,sz[N];
int hd[N],pre[M],to[M],num;
void adde(int x,int y)
{
num++;pre[num]=hd[x];hd[x]=num;to[num]=y;
}
void dfs(int v,int fa)
{
f[v]=1,g[v]=K;
for(int i=hd[v];i;i=pre[i])
{
int u=to[i];
if(u==fa)continue;
dfs(u,v);
g[v]=(1ll*g[v]*f[u]+1ll*g[u]*f[v])%mod;
f[v]=1ll*f[v]*f[u]%mod;
}
f[v]=(f[v]+g[v])%mod;
}
ll p2(ll x){return x*x%mod;}
ll pw(ll x,ll p)
{
return p?p2(pw(x,p/2))*(p&1?x:1)%mod:1;
}
int main()
{
//freopen("data.in","r",stdin);
int x,y,op,p;
scanf("%d%d%d",&n,&p,&op);
for(int i=1;i<n;i++)
{
scanf("%d%d",&x,&y);
adde(x,y),adde(y,x);
}
if(p==1)
{
printf("%lld
",pw(n,n-2));
return 0;
}
P=pw(pw(p,mod-2)-1,mod-2),K=1ll*n*P%mod;
dfs(1,0);
ll ans = g[1] * p2(pw(n,mod-2)) % mod * pw(pw(P,mod-2),n) % mod;
printf("%lld
",ans*pw(p,n)%mod);
return 0;
}
Task2
类似Task1，我们同样可以得到
其中(G(V)=(n^{n-|V|-2}prod_{i=1}^{n-|V|}a_i)^2)，V为任意一棵n个点的森林的边集
带入化简得
设(K=n^2P^{-1},ans''=sum_{V}prod_{i=1}^{n-|V|}(a_i^2K))
则(ans'=n^{-4}P^nans'')
考虑(ans'')，它相当于将n个点划分为若干连通块，每一个连通块内部构成一棵树，则每一种划分方案的贡献为(prod_{i=1}^{n-|V|}a_i^{a_i-2}(a_i^2K))，其中每一个大小为i的连通块的贡献为(i^{i-2}(i^2K)=i^iK)
根据生成函数的一些性质，如果我们设(G(x)=sum_{i=1}^infty frac{i^iK}{i!}x^i)，那么(e^{G(x)})就是(ans'')的指数型生成函数，取它的第(x^n)项再乘以(n!)就可以得到(ans'')。
多项式(exp)即可。
最后的答案为(bas^nn^{-4}P^nans'')
完整的代码（注意要特判(bas=1)的情况）：
#include<bits/stdc++.h>
using namespace std;
#define N 600007
#define ll long long
const ll mod=998244353;
const int lim=2e5;
int tp;
ll n,bas;
ll p2(ll x){return x*x%mod;}
ll pw(ll x,ll p)
{
return p?p2(pw(x,p/2))*(p&1?x:1)%mod:1;
}
namespace tp0
{
struct edge
{
int x,y;
};
bool operator <(edge a,edge b)
{
return a.x<b.x||a.x==b.x&&a.y<b.y;
}
set<edge> S;
void work()
{
if(bas==1)
{
printf("%d
",1);
return ;
}
for(int i=1;i<n;i++)
{
int x,y;
scanf("%d%d",&x,&y);
if(x>y)swap(x,y);
S.insert({x,y});
}
int cnt=0;
for(int i=1;i<n;i++)
{
int x,y;
scanf("%d%d",&x,&y);
if(x>y)swap(x,y);
if(S.find({x,y})!=S.end())cnt++;
}
printf("%lld
",pw(bas,n-cnt));
}
}
namespace tp1
{
int hd[N],pre[N],to[N],num;
ll f[N],g[N],K,P,Q;
void adde(int x,int y)
{
num++;pre[num]=hd[x];hd[x]=num;to[num]=y;
}
void dfs(int v,int fa)
{
f[v]=1,g[v]=K;
for(int i=hd[v];i;i=pre[i])
{
int u=to[i];
if(u==fa)continue;
dfs(u,v);
g[v]=(g[v]*f[u]+f[v]*g[u])%mod;
f[v]=f[v]*f[u]%mod;
}
f[v]=(f[v]+g[v])%mod;
}
void work()
{
if(bas==1)
{
printf("%lld
",pw(n,n-2));
return ;
}
P=(pw(bas,mod-2)-1+mod)%mod;
K=n*pw(P,mod-2)%mod;
Q=pw(n,2*(mod-2))*pw(P,n)%mod;
int x,y;
for(int i=1;i<n;i++)
{
scanf("%d%d",&x,&y);
adde(x,y),adde(y,x);
}
dfs(1,0);
ll ans=g[1];
ans=ans*Q%mod;
ans=ans*pw(bas,n)%mod;
printf("%lld
",ans);
}
}
namespace tp2
{
ll inv[N],fac[N],ifac[N];
int rev[N],len;
void getlen(int n)
{
for(len=1;len<=n;len<<=1);
for(int i=0;i<len;i++)
rev[i]=rev[i>>1]>>1|(i&1?len>>1:0);
}
void NTT(ll *a,int op)
{
for(int i=0;i<len;i++)
if(rev[i]<i)swap(a[rev[i]],a[i]);
for(int i=1;i<len;i<<=1)
{
ll nw=pw(3,(mod-1)/(i<<1));
for(int j=0;j<len;j+=i<<1)
{
ll w=1;
for(int k=j;k<j+i;k++)
{
ll x=a[k],y=a[k+i]*w%mod;
a[k]=(x+y)%mod,a[k+i]=(x-y+mod)%mod;
w=w*nw%mod;
}
}
}
if(op<0)
{
reverse(a+1,a+len);
ll Inv=pw(len,mod-2);
for(int i=0;i<len;i++)
a[i]=a[i]*Inv%mod;
}
}
void copy(ll *a,ll *b,int n=len)
{
for(int i=0;i<n;i++)a[i]=b[i];
for(int i=n;i<len;i++)a[i]=0;
}
ll mul_c[N],mul_d[N];
void mul(ll *t,ll *a,ll *b)
{
ll *c=mul_c,*d=mul_d;
copy(c,a),copy(d,b);
NTT(c,1),NTT(d,1);
for(int i=0;i<len;i++)c[i]=c[i]*d[i]%mod;
NTT(c,-1);
copy(t,c);
}
ll inv_c[N];
void getinv(int p,ll *a,ll *b)
{
if(p==1)return a[0]=pw(b[0],mod-2),(void)1;
getinv((p+1)/2,a,b);
getlen(2*p);
ll *c=inv_c;
copy(c,b,p);
NTT(a,1),NTT(c,1);
for(int i=0;i<len;i++)a[i]=a[i]*(2-a[i]*c[i]%mod+mod)%mod;
NTT(a,-1);
for(int i=p;i<len;i++)a[i]=0;
}
void devir(ll *a,int n)
{
for(int i=1;i<=n;i++)a[i-1]=a[i]*i%mod;
a[n]=0;
}
void inter(ll *a,int n)
{
for(int i=n;i>=0;i--)a[i+1]=a[i]*inv[i+1]%mod;
a[0]=0;
}
ll ln_c[N];
void getln(int n,ll *a,ll *b)
{
ll *c=ln_c;
getlen(2*n);
copy(c,b,n);
getinv(n,a,c);
devir(c,n);
getlen(2*n);
mul(a,a,c);
inter(a,n);
for(int i=n;i<len;i++)a[i]=0;
}
ll exp_c[N];
void getexp(int p,ll *a,ll *b)
{
if(p==1)return a[0]=1,(void)1;
getexp((p+1)/2,a,b);
getlen(2*p);
ll *c=exp_c;
copy(c,a,0);
getln(p,c,a);
getlen(2*p);
for(int i=0;i<p;i++)c[i]=(b[i]-c[i]+mod)%mod;
c[0]=(c[0]+1)%mod;
mul(a,a,c);
for(int i=p;i<len;i++)a[i]=0;
}
void Init()
{
fac[0]=1;
for(int i=1;i<=lim;i++)fac[i]=fac[i-1]*i%mod;
ifac[lim]=pw(fac[lim],mod-2);
for(int i=lim;i>=1;i--)ifac[i-1]=ifac[i]*i%mod;
inv[1]=1;
for(int i=2;i<=lim;i++)
inv[i]=mod-mod/i*inv[mod%i]%mod;
}
ll f[N],g[N];
void work()
{
Init();
ll P,K;
if(bas==1)
{
printf("%lld
",pw(n,2*(n-2)));
return ;
}
P=(pw(bas,mod-2)-1+mod)%mod;
K=n*n%mod*pw(P,mod-2)%mod;
for(int i=1;i<=n;i++)
f[i]=pw(i,i)*K%mod*ifac[i]%mod;
getexp(n+1,g,f);
ll ans=g[n]*fac[n]%mod;
ans=ans*pw(n,4*(mod-2))%mod*pw(P,n)%mod;
ans=ans*pw(bas,n)%mod;
printf("%lld
",ans);
}
}
int main()
{
scanf("%lld%lld%d",&n,&bas,&tp);
if(tp==0)tp0::work();
else if(tp==1)tp1::work();
else tp2::work();
return 0;
}