[BZOJ2152] 聪聪可可

Description

给定一棵树，树上每条边有权值，两人每次从树上随机选择两个点（可以重复），求两点路径上长为 (3) 的倍数的概率。

Solution

每层分治到重心 (p) 时计算所有经过 (p) 的路径，这里我们先只算不重复的组合，最后乘 (2) 再加上两点重合的情况

对重心 (p) 和当前枚举的子树 (q)，分别统计 (p) 发出的对应和模的路径的条数以及 (q) 发出的对应点权和模的路径的条数

具体过程中，每次 DFS 求出 (q) 的 map，计算答案，然后并入 (p) 的 map 即可

#include <bits/stdc++.h>
using namespace std;
#define int long long
#define foradj(p) for(pair<int,int> pr:g[p])
#define unpii(q,w) int q=pr.first, w=pr.second;
const int N = 1000005;

vector <pair<int,int> > g[N];
int dis[N],siz[N],n,m,t1,t2,t3,k[N],ans,vis[N],sum,rt,rtval;
vector<int> st,wl;
set<int> s;
map<int,int> mp_main,mp_sub;

void findroot(int p,int ff)
{
    siz[p]=1;
    int mx=0;
    foradj(p)
    {
        unpii(q,w);
        if(!vis[q] && q!=ff)
        {
            findroot(q,p);
            siz[p]+=siz[q];
            mx=max(mx,siz[q]);
        }
    }
    mx=max(mx,sum-siz[p]);
    if(mx<rtval) rtval=mx, rt=p;
}

void dfs(int p,int ff)
{
    mp_sub[dis[p]%3]++;
    foradj(p)
    {
        unpii(q,w);
        if(!vis[q] && q!=ff)
        {
            dis[q]=dis[p]+w;
            dis[q]%=3;
            dfs(q,p);
        }
    }
}

void solve(int p)
{
    vis[p]=1;
    mp_main.clear();
    mp_main[0]++;
    foradj(p)
    {
        unpii(q,w);
        if(!vis[q])
        {
            dis[q]=w%3;
            dfs(q,p);
            ans+=mp_main[0]*mp_sub[0]+mp_main[1]*mp_sub[2]+mp_main[2]*mp_sub[1];
            for(auto i:mp_sub) mp_main[i.first]+=i.second;
            mp_sub.clear();
        }
    }
    foradj(p)
    {
        unpii(q,w);
        if(!vis[q])
        {
            sum=siz[q];
            rtval=1e18;
            findroot(q,0);
            solve(rt);
        }
    }
}

signed main()
{
    ios::sync_with_stdio(false);
    cin>>n;
    for(int i=1;i<n;i++)
    {
        cin>>t1>>t2>>t3;
        g[t1].push_back({t2,t3});
        g[t2].push_back({t1,t3});
    }
    sum=n;
    rtval=1e18;
    findroot(1,0);
    solve(rt);
    ans*=2;
    ans+=n;
    {
        int p=ans,q=n*n;
        int g=__gcd(p,q);
        if(g>1)
        {
            p/=g;
            q/=g;
        }
        cout<<p<<"/"<<q<<endl;
    }
}