题目链接：Codeforces 1156D

你有一棵由 $n$ 个节点组成的树，每条边上有一个数字 $0$ 或 $1$。

假如我们从节点 $x$ 开始沿着简单路径到达节点 $y$，一旦经过了数字为 $1$ 的边就不会经过数字为 $0$ 的边。那么我们称二元组 $(x, y)$（$x \ne y$）是合法的。

你的任务是找出树上合法的二元组数量。

数据范围：$2 \le n \le 2 \times 10 ^ 5$。

Solution

我们对这条路径分为 $3$ 种情况讨论，对这些情况分别计算贡献。

1. 只包含 $0$ 边。
2. 只包含 $1$ 边。
3. 起初为 $0$ 边，后来为 $1$ 边。

对于 $0, 1$ 边分别建立其生成树森林，设数字为 $j$ 的第 $i$ 棵树的大小为 $size(i, j)$，那么对于三种情况分别有：

1. 贡献为 $\sum A_{size(i, 0)}^{2}$（每个集合产生一次贡献）。
2. 贡献为 $\sum A_{size(i, 1)}^{2}$（每个集合产生一次贡献）。
3. 考虑枚举交界的那个点，那么在其所在 $0, 1$ 边生成树中各选一个点，贡献为 $(size(belong(i), 0) - 1) \cdot (size(belong(i), 1) - 1)$。其中 $belong(i)$ 表示点 $i$ 所在集合，当然因数字而异（每个点产生一次贡献）。

当然也可以通过树形 $\text{DP}$ 在 $\mathcal O(n)$ 的更优复杂度内求解。但是鉴于码量和细节问题，笔者选择并查集的解法。

时间复杂度：$\mathcal O(n \log n)$。

Code

#include <cstdio>

const int N = 2e5 + 5;

int n, u[N], v[N], w[N], fa[N][2], sz[N][2];

int find(int x, int c) {
    return fa[x][c] == x ? x : fa[x][c] = find(fa[x][c], c);
}
void merge(int x, int y, int c) {
    x = find(x, c), y = find(y, c);
    if (x == y) return;
    fa[y][c] = x, sz[x][c] += sz[y][c];
}
int main() {
    scanf("%d", &n);
    for (int i = 1; i <= n; i++) {
        fa[i][0] = fa[i][1] = i;
        sz[i][0] = sz[i][1] = 1;
    }
    for (int i = 1; i < n; i++) {
        int u, v, w;
        scanf("%d%d%d", &u, &v, &w);
        merge(u, v, w);
    }
    long long ans = 0;
    for (int i = 1; i <= n; i++) {
        if (fa[i][0] == i) ans += 1LL * sz[i][0] * (sz[i][0] - 1);
        if (fa[i][1] == i) ans += 1LL * sz[i][1] * (sz[i][1] - 1);
        ans += 1LL * (sz[find(i, 0)][0] - 1) * (sz[find(i, 1)][1] - 1);
    }
    printf("%lld\n", ans);
    return 0;
}