Codeforces 1156D 带权并查集

题意：给你一颗树，树边的权值可能是0或1，问先走0边，再走1边，或者只走1边的路径有多少条？

思路：对于一个点，假设通过0边相连的点一共有x个（包括自己），通过1边相连的有y个（包括自己），那么对答案的贡献为x * y - 1，意思是以x个点为起点，以y个点为终点了路径条数，-1是因为自己到自己的不算路径，用带权并查集去统计x和y即可。

代码：

#include <bits/stdc++.h>
#define LL long long
#define db double
#define INF 0x3f3f3f3f
#define pii pair<int, int>
using namespace std;
const int maxn = 200010;
int a[maxn];
int f[2][maxn], sz[2][maxn];
int get(int pos, int x) {
	if(x == f[pos][x]) return x;
	return f[pos][x] = get(pos, f[pos][x]);
}
void merge(int pos, int x, int y) {
	int x1 = get(pos, x), y1 = get(pos, y);
	if(x1 == y1) return;
	f[pos][x1] = y1;
	sz[pos][y1] += sz[pos][x1]; 
}
int main() {
	int n, x, y, z;
	LL ans = 0;
	scanf("%d", &n);
	for (int i = 1; i <= n; i++) {
		f[0][i] = f[1][i] = i;
		sz[0][i] = sz[1][i] = 1;
	}
	for (int i = 1; i < n; i++) {
		scanf("%d%d%d", &x, &y, &z);
		merge(z, x, y);
	}
	for (int i = 1; i <= n; i++) {
		ans +=  ((long long)sz[0][get(0, i)] * sz[1][get(1, i)]) - 1;
	}
	printf("%lld\n", ans);
}