CF1770E Koxia and Tree

CF1770E Koxia and Tree

题目大意

有一棵$n$个结点的无根树，边从$1$到$n-1$标号，第$i$条边连接结点$u_i$和$v_i$。

有$k$个结点带有标记，分别是$a_1,a_2,\dots a_k(\forall i\ne j,a_i\ne a_j)$。接下来，进行如下操作：

• 给每条边随机定一个方向
• 按照编号的顺序，对于每条边$u\to v$，如果结点$u$有标记而结点$v$没有，那么就把结点$u$的标记转移到结点$v$上
• 随机选取两个不同的标记点，计算它们的距离。

请你计算第三步中距离的期望值。对$998244353$取模。

题解

我们先考虑如果不转移标记，距离的期望值为多少。

此时距离的期望值就是任意两点距离的总和减去任意选择两个点的方案数，那么只需求出任意两点距离的总和即可。我们可以分别计算每条边的贡献。因为路径上有这条边的两个点一定是一个在左一个在右，所以一条边的贡献为其左边的结点中有标记的点的个数与右边的结点中有标记的点的个数之积。令一个点为根，维护每个结点所在子树的有标记的点的个数，即可用dfs来$O(n)$求出答案。

我们再考虑会转移标记的情况下的距离期望值。

首先我们知道，在遍历一条边之前，没有标记能跨越这条边，所以这条边两边的有标记的点是不变的。对于一条边$(u,v)$，假设$u$为$v$的父亲，那么$u$这边有标记的点的个数为$k-s_v$，$v$这边有标记的点的个数$s_v$。

令$u$点有标记的概率为$f_u$，$v$点有标记的概率为$f_v$。我们可以分类讨论。

• 两端都有点，则贡献为$f_u\times f_v\times (k-s_v)\times s_v$
• $u$有点，$v$没有点，如果方向为$u\to v$，则$u$点的标记转移到$v$，否则不转移，概率各占$\frac{1}{2}$，贡献为
  $(1-f_u)\times f_v\times \frac{(k-s_v-1)\times (s_v+1)+(k-s_v)\times s_v}{2}$
• $u$没有点，$v$有点，如果方向为$v\to u$，则$v$点的标记转移到$u$，否则不转移，概率各占$\frac{1}{2}$，贡献为
  $f_u\times (1-f_v)\times \frac{(k-s_v+1)\times (s_v-1)+(k-s_v)\times s_v}{2}$
• 两端都没有点，贡献为$(1-f_u)\times (1-f_v)\times (k-s_v)\times s_v$

这四个式子之和就是这条边对答案的贡献，累加即可。

那么遍历这条边后，两边有点的概率会变为多少呢？

• 两边都有点，遍历这条边后$u$有点和$v$有点的概率为$f_u\times f_v$
• $u$有点，$v$没有点，遍历这条边后，两种方向的概率各占$\frac{1}{2}$，$u$有点和$v$有点的概率为$\frac{f_u\times (1-f_v)}{2}$
• $v$有点，$u$没有点，遍历这条边后，两种方向的概率各占$\frac{1}{2}$，$u$有点和$v$有点的概率为$\frac{(1-f_u)\times f_v}{2}$
• 两边没有点，遍历这条边后$u$有点和$v$有点的概率为$(1-f_u)\times (1-f_v)$

那么遍历这条边后$f_u=f_v=\frac{f_u+f_v}{2}$。

最后的答案要乘上方案数的逆元。

时间复杂度为$O(n)$。

code

#include<bits/stdc++.h>
using namespace std;
int n,x,y,tot=1,fa[300005],d[600005],l[600005],r[600005];
long long k,ans=0,siz[300005],p[300005];
long long ny2=499122177,mod=998244353;
void add(int xx,int yy){
	l[++tot]=r[xx];d[tot]=yy;r[xx]=tot;
}
void dfs(int u,int f){
	fa[u]=f;siz[u]=p[u];
	for(int i=r[u];i;i=l[i]){
		if(d[i]==f) continue;
		dfs(d[i],u);
		siz[u]+=siz[d[i]];
	}
}
long long mi(long long t,long long v){
	if(!v) return 1;
	long long re=mi(t,v/2);
	re=re*re%mod;
	if(v&1) re=re*t%mod;
	return re;
}
int main()
{
	scanf("%d%lld",&n,&k);
	for(int i=1;i<=k;i++){
		scanf("%d",&x);p[x]=1;
	}
	for(int i=1;i<n;i++){
		scanf("%d%d",&x,&y);
		add(x,y);add(y,x);
	}
	dfs(1,0);
	for(int i=2;i<=tot;i+=2){
		int u=d[i^1],v=d[i];
		if(fa[u]==v) swap(u,v);
		ans=(ans+p[u]*p[v]%mod*(k-siz[v])%mod*siz[v]%mod
			+(mod+1-p[u])*(mod+1-p[v])%mod*(k-siz[v])%mod*siz[v]%mod
			+p[u]*(mod+1-p[v])%mod*((k-siz[v])*siz[v]%mod+(k-siz[v]-1)*(siz[v]+1)%mod)%mod*ny2%mod
			+(mod+1-p[u])*p[v]%mod*((k-siz[v])*siz[v]%mod+(k-siz[v]+1)*(siz[v]-1)%mod)%mod*ny2%mod)%mod;
		p[u]=p[v]=(p[u]+p[v])*ny2%mod;
	}
	printf("%lld",ans*mi(k*(k-1)/2%mod,mod-2)%mod);
	return 0;
}