思路:经过的路径可以分为三类,只有0,只有1,先0后1,可以用并查集将相邻同数边合并,如果从x到y的路径包含两种类型的边,则存在一个顶点v,这样从x到v的路径只经过0边,而从v到y的路径只经过1边。那么,让我们迭代这个顶点v,在图选择一个顶点作为x,选择另一个顶点作为y,并将选择它们的方法的数量加在答案中。
#include<bits/stdc++.h>
using namespace std;
const int N = 2e5 + 7;
typedef long long ll;
int n, x, y, z;
ll ans;
struct node {
int fa[N], num[N];
void init() {
for(int i = 0; i <= n; i++)
fa[i] = i, num[i] = 1;
}
int F(int x) {
return fa[x] == x ? x : fa[x] = F(fa[x]);
}
} no[2];
int main() {
cin >> n;
no[0].init(), no[1].init();
for(int i = 1; i < n; i++) {
cin >> x >> y >> z;
x = no[z].F(x), y = no[z].F(y);
if(x != y) {
no[z].fa[x] = y;
no[z].num[y] += no[z].num[x];
}
}
for(int i = 1; i <= n; i++)
ans += 1ll * no[0].num[no[0].F(i)] * no[1].num[no[1].F(i)] - 1;
cout << ans;
return 0;
}