2020 China Collegiate Programming Contest, Weihai Site C. Rencontre(lca+树形dp)

Posted 2021-05-25 issue是fw

LINK

题意

给定一颗$n$个节点的树

给定集合$a,b,c$,从每个集合随机选一个点出来

从树中找一个点$x$使得三点到$x$的距离和最小

求这个距离和的期望

---

本质就是求所有组合的距离和大小

对于给定的三点$x,y,z$,显然这个距离和是

$\frac{(dis(x,y)+dis(y,z)+dis(x,z))}{2}$

又因为$$dis(x,y)=deep[x]+deep[y]-2\ast deep[lca(x,y)]$$

所以我们可以枚举集合$a,b$,集合$a,c$,集合$b,c$两两计算贡献

比如现在是算$a,b$集合点间的贡献,从$a$中拿出点$x$,从$b$中拿出点$y$

那么$deep[x]$被计算了$b.size()$次,$deep[y]$被计算了$a.size()$

然后我们还要计算每个点被作为$lca$多少次,设$k$点作为$lca$共$siz$次

那么答案还要减去$2\ast deep[k]\ast siz$

但是注意,现在的答案是$a,b$集合间任意两点的距离和

这个距离和还需要对每个$c$集合的点都会作用一次,所以乘以$c.size()$

同理对集合$b,c$,对集合$a,c$也计算一遍

最后算期望,除以总方案数即可

#include <bits/stdc++.h>
using namespace std;
const int maxn = 1e6+10;
vector<int>vec[4];
struct edge{
	int to,nxt,w;
}d[maxn]; int head[maxn],cnt=1;
void add(int u,int v,int w)
{
	d[++cnt] = ( edge ){v,head[u],w},head[u] = cnt;
}
int n,deep[maxn],siz[maxn][3];
double sum=0; 
void PRE(int u,int fa)
{
	for(int i=head[u];i;i=d[i].nxt )
	{
		int v =  d[i].to;
		if( v==fa )	continue;
		deep[v] = deep[u]+d[i].w;
		PRE(v,u);
	}
}
void dfs(int u,int fa)
{
	if( siz[u][0] && siz[u][1] )	sum -= 2*deep[u];
	for(int i=head[u];i;i=d[i].nxt )
	{
		int v = d[i].to;
		if( v==fa )	continue;
		dfs(v,u);
		sum -= 1ll*2*siz[u][0]*siz[v][1]*deep[u];
		sum -= 1ll*2*siz[u][1]*siz[v][0]*deep[u];
		siz[u][0] += siz[v][0], siz[u][1] += siz[v][1];
	}
}
double solve( vector<int>a,vector<int>b,vector<int>c)
{
	for(auto v:a )	siz[v][0]++;
	for(auto v:b )	siz[v][1]++;
	sum = 0; dfs(1,1);
	for(auto v:a )	sum += 1ll*deep[v]*b.size();
	for(auto v:b )	sum += 1ll*deep[v]*a.size();
	sum = sum*c.size();
	for(int i=1;i<=n;i++)	siz[i][0] = siz[i][1] = 0;
	return sum;
}
int main()
{
	cin >> n;
	for(int i=1;i<n;i++)
	{
		int l,r,w; scanf("%d%d%d",&l,&r,&w);
		add(l,r,w); add(r,l,w);
	}
	for(int i=1;i<=3;i++)	
	{
		int k; cin >> k;
		for(int j=1;j<=k;j++)
		{
			int x; scanf("%d",&x);
			vec[i].push_back( x );
		}
	}
	PRE(1,1);
	double ans = 0;
	ans += solve( vec[1],vec[2],vec[3] );
	ans += solve( vec[1],vec[3],vec[2] );
	ans += solve( vec[2],vec[3],vec[1] );
	ans /= 2;
	double ANS = 1.0*ans/( 1ll*vec[1].size()*vec[2].size()*vec[3].size() );
	printf("%.19lf",ANS);
}