小K的疑惑

题目

给定一个有 𝑁 个点的树，每条边的长度有一个边权，现在定义 𝑑𝑖𝑠(𝑖,𝑗) 代表第 𝑖 个点到第 𝑗 个点的距离模 2 之后的结果。
问有多少 (𝑖,𝑗,𝑘) 满足，𝑑𝑖𝑠(𝑖,𝑗) = 𝑑𝑖𝑠(𝑗,𝑘) = 𝑑𝑖𝑠(𝑖,𝑘)。

解题思路

可以将两点的距离转换成到根节点的距离。可以任意选一个节点作为根节点，比如节点 1。 要么是 0，要么是 1。

同理，
要使 ，则要 。
同理，要使 ，则要 。
综上，当 时，由 。

使用 DFS 算法计算从根节点 1 到每个节点的距离，并统计每种距离的数目。

表示 的节点数。

根据数学的排列，得出答案为 。

C++代码

#include<iostream>
#include<vector>
using namespace std;

vector<vector<vector<int>>> edges;
vector<bool> vis;
int cnt[3];

void dfs(int s, int d){
    for(auto e : edges[s]){
        if(!vis[e[0]]){
            vis[e[0]] = true;
            int tmp = (d + e[1]) % 2;
            cnt[tmp] += 1;
            dfs(e[0], tmp);
        }
    }
}

int main(){
    int n;
    cin >> n;
    edges.resize(n+1);
    int s, e, d;
    for(int i=0; i<n; ++i){
        cin >> s >> e >> d;
        d %= 2;
        edges[s].push_back({e,d});
        edges[e].push_back({s,d});
    }
    vis.resize(n+1, false);
    vis[1] = true;
    cnt[0] += 1;    // 根节点本身与自己的距离是 0
    dfs(1,0);
    long long ans = 1LL*cnt[0]*cnt[0]*cnt[0] + 1LL*cnt[1]*cnt[1]*cnt[1];
    cout << ans << endl;
    return 0;
}