洛谷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;
}