题目描述
给定一棵\(n\)个点的树，在\(k\)个位置上存在蝴蝶，我们需要给\(n-1\)条边定向，如果一条边的起点有蝴蝶且终点没有蝴蝶，那么蝴蝶将被移动到终点，我们会按照给定边的顺序移动，问最终所有蝴蝶的树上距离的和的期望，答案除于\(\frac{k(k-1)}{2}\)，对\(998244353\)取模
\[k\le n\le 300000 \]题解
首先考虑第一个问题，我们怎么求\(k\)个点的距离和，可以随便钦定一个点为根，求出\(siz_u\)表示子树\(u\)中的指定结点个数，答案显然就是\(\sum siz_u(k-siz_u)\)
意识到每条边只会经历一遍，所以\(siz_u\)只会变化\(1\)，可以通过\(u\)到父亲的边的移动情况讨论得出
由于\(u\)到其父亲上的边未走过之前，\(u\)和其父亲有蝴蝶的概率是互不影响的，所以我们只需要知道\(u\)和其父亲上有蝴蝶的概率就可以算出\(\Delta ans\)，新的概率也可以动态维护\(p_u=p_{fa_u}=\frac{p_u+p_{fa_u}}{2}\)
\(\text{code}\)
#include<cstdio>
#include<algorithm>
#include<vector>
#define ll long long
using namespace std;
const ll mod=998244353,inv2=499122177;
ll ksm(ll a,ll b)
{
if(b==0) return 1;
ll tmp=ksm(a,b>>1);
if(b&1) return tmp*tmp%mod*a%mod;
else return tmp*tmp%mod;
}
void read(int &res)
{
res=0;char ch=getchar();
while(ch<'0'||ch>'9') ch=getchar();
while('0'<=ch&&ch<='9') res=(res<<1)+(res<<3)+(ch^48),ch=getchar();
}
const int N=3e5+100;
int n,k,a[N+10],A[N+10],B[N+10];
vector<int> e[N+10];
bool vis[N+10];
int fa[N+10],siz[N+10];
ll ans,p[N+10];
void dfs(int u)
{
siz[u]=vis[u];
for(int v:e[u])
{
if(v==fa[u]) continue;
fa[v]=u,dfs(v),siz[u]+=siz[v];
}
ans+=1ll*siz[u]*(k-siz[u]);
}
int main()
{
read(n),read(k);
for(int i=1;i<=k;i++) read(a[i]),vis[a[i]]=true,p[a[i]]=1;
for(int i=1;i<n;i++) read(A[i]),read(B[i]),e[A[i]].push_back(B[i]),e[B[i]].push_back(A[i]);
dfs(1);
for(int i=1;i<n;i++)
{
if(fa[A[i]]==B[i]) swap(A[i],B[i]);
ans+=p[A[i]]*(mod+1-p[B[i]])%mod*(1ll*(siz[B[i]]+1)*(k-siz[B[i]]-1)%mod+mod-1ll*siz[B[i]]*(k-siz[B[i]])%mod)%mod*inv2%mod,ans%=mod;
ans+=p[B[i]]*(mod+1-p[A[i]])%mod*(1ll*(siz[B[i]]-1)*(k-siz[B[i]]+1)%mod+mod-1ll*siz[B[i]]*(k-siz[B[i]])%mod)%mod*inv2%mod,ans%=mod;
p[A[i]]=p[B[i]]=(p[A[i]]+p[B[i]])*inv2%mod;
}
printf("%lld\n",ans*ksm(1ll*k*(k-1)%mod*inv2%mod,mod-2)%mod);
return 0;
}