BZOJ2152 聪聪可可

【题意】

求树上路径长度为3的倍数的路径数

【分析】

对于每个root，记录长度除以3余0/1/2的数量，计算答案即可

【代码】

#include<bits/stdc++.h>
using namespace std;
const int maxn=2e4+5;
const int inf=0x3f3f3f3f;
int n,head[maxn],cnt;
struct edge
{
    int to,nxt,v;    
}e[maxn<<1];
void add(int x,int y,int z)
{
    e[++cnt].nxt=head[x]; e[cnt].to=y; e[cnt].v=z; head[x]=cnt;
    e[++cnt].nxt=head[y]; e[cnt].to=x; e[cnt].v=z; head[y]=cnt;
}
int root,dis[maxn],siz[maxn],vis[maxn],f[maxn],sum;
void getroot(int x,int fa)
{
    siz[x]=1; f[x]=0;
    for(int i=head[x];i;i=e[i].nxt)
    {
        int to=e[i].to;
        if(to==fa || vis[to]) continue;
        getroot(to,x);
        siz[x]+=siz[to];
        f[x]=max(f[x],siz[to]);
    }
    f[x]=max(f[x],sum-siz[x]);
    root=(f[x]<f[root])?x:root;
}
int t[4];
void getdeep(int x,int fa)
{
    t[dis[x]]++;
    for(int i=head[x];i;i=e[i].nxt)
    {
        int to=e[i].to;
        if(to==fa || vis[to]) continue;
        dis[to]=(dis[x]+e[i].v)%3;
        getdeep(to,x);
    }
}
int ans;
int calc(int x,int len)
{
    t[0]=t[1]=t[2]=0;
    dis[x]=len;
    getdeep(x,0);
    return t[1]*t[2]*2+t[0]*t[0];
}
void solve(int x)
{
    ans+=calc(x,0);
    vis[x]=1;
    for(int i=head[x];i;i=e[i].nxt)
    {
        int to=e[i].to;
        if(vis[to]) continue;
        ans-=calc(to,e[i].v);
        sum=siz[to];
        root=0;
        getroot(to,x);
        solve(root);
    }
}
int gcd(int x,int y)
{
    return y==0?x:gcd(y,x%y);
}
int main()
{
    //freopen("a.in","r",stdin);
    //freopen("a.out","w",stdout);
    scanf("%d",&n);
    for(int i=1;i<n;i++)
    {
        int x,y,z;
        scanf("%d%d%d",&x,&y,&z); z%=3;
        add(x,y,z);
    }
    sum=n; f[0]=inf;
    getroot(1,0);
    solve(root);
    int g=gcd(ans,n*n);
    printf("%d/%d\n",ans/g,n*n/g);
    return 0;
}