牛客算法周周练18-A 小K的疑惑

题目链接：https://ac.nowcoder.com/acm/contest/6760/A
题意：树上找出满足dis(i,j)=dis(j,k)=dis(i,k)的三元组的个数。且dis(i,j)的定义是i到j的距离对2取模
题解：我们可以将两个点之间的距离，全部转换成到根的距离。我的做法就是再找一个节点为0，使它向1连一条长度为0的点，并以0为根（以1为根其实一样），然后我们讲dis(i,j)看成(dis(0,i)+dis(0,j))%2因为是对2取模，所以lca(i,j)到0的路径长会计算两次，取模后没有影响，那么我们稍微推导一下发现，i,j,k三元组若满足题目条件，则i,j,k到根的距离奇偶性相同。这里就考虑全为偶，假设偶数个节点数为x,那么对于i有x种可能，j有x种可能，k也有x种可能，故ans=xxx+yyy;

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
int n,x,y,z;
const int maxn=10005;
struct node{int v,w,nxt;}g[maxn<<1];
int head[maxn],tot;
void add_edge(int u,int v,int w){g[++tot]={v,w,head[u]};head[u]=tot;}
int d[maxn];
int o,j;
void build(int u,int fa,int now){
    d[u]=now;
    if (u!=0){
    if (d[u]&1)j++;
    else o++;
    }
    for (int i=head[u];i;i=g[i].nxt){
        int v=g[i].v;
        if (v!=fa){
            build(v,u,(now+g[i].w)%2);
        }
    }
}
int main(){
    cin>>n;
    for (int i=1;i<n;i++){
        scanf("%d%d%d",&x,&y,&z);
        add_edge(x,y,z);
        add_edge(y,x,z);
    }
    add_edge(0,1,0);
    build(0,-1,0);
    ll ans=0;
    ans=1ll*j*j*j+1ll*o*o*o;
    cout<<ans<<endl;
    return 0;
}