洛谷P2634 [国家集训队]聪聪可可 点分治模板

题意

在一棵树上任意选两个点，求它们距离模3为0的概率。

分析

树分治模板

Code

#include<bits/stdc++.h>
#define fi first
#define se second
#define pb push_back
#define lson l,mid,p<<1
#define rson mid+1,r,p<<1|1
#define ll long long
using namespace std;
const int inf=1e9;
const int mod=1e9+7;
const int maxn=2e4+10;
typedef pair<int,int> pii;
vector<pii>g[maxn];
int n,sz[maxn],vis[maxn],maxp[maxn],f[3],c[3],ans,rt,sum;
void getrt(int u,int fa){
	sz[u]=1;maxp[u]=0;
	for(pii x:g[u]){
		if(x.fi==fa||vis[x.fi]) continue;
		getrt(x.fi,u);
		sz[u]+=sz[x.fi];
		maxp[u]=max(maxp[u],sz[x.fi]);
	}
	maxp[u]=max(maxp[u],sum-sz[u]);
	if(maxp[u]<maxp[rt]) rt=u;
}
void getdis(int u,int fa,int d){
	f[d%3]++;
	for(pii x:g[u]){
		if(x.fi==fa||vis[x.fi]) continue;
		getdis(x.fi,u,d+x.se);
	}
}
void calc(int u){
	for(pii x:g[u]){
		if(vis[x.fi]) continue;
		f[0]=f[1]=f[2]=0;
		getdis(x.fi,u,x.se);
		ans+=f[0]*c[0]*2+f[1]*c[2]*2+f[2]*c[1]*2;
		c[0]+=f[0];c[1]+=f[1];c[2]+=f[2];
	}
	c[0]=1;c[1]=c[2]=0;
}
void solve(int u){
	vis[u]=c[0]=1;calc(u);
	for(pii x:g[u]){
		if(vis[x.fi]) continue;
		sum=sz[x.fi],maxp[rt=0]=inf;
		getrt(x.fi,0);solve(rt);
	}
}
ll gcd(ll a,ll b){
	return b==0?a:gcd(b,a%b);
}
int main(){
	ios::sync_with_stdio(false);
	//freopen("in","r",stdin);
	cin>>n;
	for(int i=1,x,y,w;i<n;i++){
		cin>>x>>y>>w;
		g[x].pb(pii(y,w));g[y].pb(pii(x,w));
	}
	sum=maxp[rt]=n;
	getrt(1,0);solve(rt);
	int a=ans+n,b=n*n;
	int t=gcd(a,b);
	a/=t;b/=t;
	cout<<a<<'/'<<b<<endl;
	return 0;
}