题目链接

题面：

题意：
求树上两点之间距离为3的倍数的点对数量。

代码：

#include<iostream>
#include<cstdlib>
#include<cstdio>
#include<algorithm>
#include<cstring>
#include<string>
#include<cmath>
#include<queue>
#include<map>
#include<vector>
#define ll long long
#define llu unsigned ll
using namespace std;
const int maxn=41000;
const int inf=0x3f3f3f3f;
int head[maxn],ver[maxn],edge[maxn],nt[maxn];
int max_part[maxn],si[maxn],cnt[3];
bool ha[maxn];
int tot=1,root,max_size,ans=0,n,k;

ll gcd(ll a,ll b)
{
    if(b==0) return a;
    return gcd(b,a%b);
}

void add(int x,int y,int z)
{
    ver[++tot]=y,edge[tot]=z;
    nt[tot]=head[x],head[x]=tot;
}

void dfs_size(int x,int fa)
{
    si[x]=1,max_part[x]=0;
    for(int i=head[x];i;i=nt[i])
    {
        int y=ver[i];
        if(ha[y]||y==fa) continue;
        dfs_size(y,x);
        si[x]+=si[y];
        max_part[x]=max(max_part[x],si[y]);
    }
}

void dfs_root(int now_root,int x,int fa)
{
    max_part[x]=max(max_part[x],si[now_root]-si[x]);

    if(max_size>max_part[x])
    {
        max_size=max_part[x];
        root=x;
    }
    for(int i=head[x];i;i=nt[i])
    {
        int y=ver[i];
        if(ha[y]||y==fa) continue;
        dfs_root(now_root,y,x);
    }
}

void dfs_dis(int x,int fa,int dis)
{
    cnt[dis]++;
    for(int i=head[x];i;i=nt[i])
    {
        int y=ver[i],z=edge[i];
        if(ha[y]||y==fa) continue;
        dfs_dis(y,x,(dis+z)%3);
    }
}

int get_num(int x,int d)
{
    cnt[0]=cnt[1]=cnt[2]=0;
    dfs_dis(x,-1,d%3);
    return cnt[0]*cnt[0]+cnt[1]*cnt[2]*2;
}

void dfs(int x)
{
    max_size=inf;
    dfs_size(x,-1);
    dfs_root(x,x,-1);
    ans+=get_num(root,0);
    ha[root]=1;
    for(int i=head[root];i;i=nt[i])
    {
        int y=ver[i],z=edge[i];
        if(ha[y]) continue;
        ans-=get_num(y,z);
        dfs(y);
    }
}

int main(void)
{
    scanf("%d",&n);
    int x,y,z;
    for(int i=1;i<n;i++)
    {
        scanf("%d%d%d",&x,&y,&z);
        add(x,y,z);add(y,x,z);
    }
    dfs(1);
    int t=gcd(n*n,ans);
    printf("%d/%d\n",ans/t,n*n/t);

    return 0;
}