https://codeforces.com/contest/1156/problem/D
乱搞
只连接0边搞成一个图,只连接1边搞成一个图。
两个图中的联通块内每个点和别的点可以统计答案,贡献为
c
n
t
∗
(
c
n
t
−
1
)
cnt * (cnt - 1)
cnt∗(cnt−1)(
c
n
t
cnt
cnt是联通块内结点数量)
然后先0再1的情况去枚举0到1的那个中间点
v
v
v,用
q
q
q 表示
v
v
v 所在0块的节点数,用
p
p
p 表示
v
v
v所在1块的节点数,则该点对答案的贡献为
(
p
−
1
)
∗
(
q
−
1
)
(p - 1)*(q - 1)
(p−1)∗(q−1)。
#include <cstdio>
#include <cstring>
#include <iostream>
#include <algorithm>
using namespace std;
typedef long long ll;
const int Maxn = 400100;
struct node {
int next, to;
} e[2][Maxn];
int h[2][Maxn], tot[2], n; ll ans;
int bel[2][Maxn], sum[2][Maxn], cnt[2];
void Add(int u, int v, int c) {
e[c][++tot[c]].to = v;
e[c][tot[c]].next = h[c][u];
h[c][u] = tot[c];
}
void Dfs(int u, int c, int p) {
bel[p][u] = c; ++sum[p][c];
for(int i = h[p][u]; i; i = e[p][i].next) {
int v = e[p][i].to;
if(!bel[p][v]) {
Dfs(v, c, p);
}
}
}
int main() {
scanf("%d", &n);
for(int i = 1; i < n; ++i) {
int x, y, c;
scanf("%d%d%d", &x, &y, &c);
Add(x, y, c); Add(y, x, c);
}
for(int p = 0; p <= 1; ++p) {
for(int i = 1; i <= n; ++i) {
if(!bel[p][i]) {
Dfs(i, ++cnt[p], p);
}
}
}
for(int p = 0; p <= 1; ++p) {
for(int i = 1; i <= cnt[p]; ++i) {
ans += (ll)sum[p][i] * (ll)(sum[p][i] - 1);
}
}
// cout << ans << endl;
for(int i = 1; i <= n; ++i) {
ans += (ll)(sum[0][bel[0][i]] - 1) * (ll)(sum[1][bel[1][i]] - 1);
}
cout << ans << endl;
return 0;
}