Codeforces 1156D 0-1-Tree（并查集）

传送门

题意：给一棵n个点，边被黑白染色的树，问有多少对点对满足从x沿最短路径走到y，经过白边(1-edges)后不经过黑边(0-edges)。

题解：开两个并查集，第一个只连黑边，第二个只连白边，讨论三种路径：①黑②白③黑转白。前两种在并查集merge的时候可以计算，第三种路径怎么搞？对于某个点x，经过它的第三种路径有条。为什么这么不会算重，因为树是没有环的（脑补一下似乎没毛病）。

做完发现可以不用讨论，直接。减去那个1就是去掉(pos, pos)这条路径。

WA了好久发现应该find的地方写成了直接调用fa数组......

P.S. 好像也可以树形dp搞

#include<cstdio>
#include<cstring>
#include<iostream>
#include<algorithm>
using namespace std;
typedef long long ll;
const int N=2e5+4; 
int n;
int fa[2][N],siz[2][N];
ll ans;
struct Edge {
	int u,v,c;
}e[N];
inline int read() {
	int x=0,f=1;char c=getchar();
	while (c<'0'||c>'9') {if (c=='-') f=-1;c=getchar();}
	while (c>='0'&&c<='9') x=x*10+c-'0',c=getchar();
	return x*f;
}
inline int find(int x,int c) {
	return fa[c][x]==x?x:fa[c][x]=find(fa[c][x],c);
}
inline void merge(int u,int v,int c) {
	int fu=find(v,c),fv=find(u,c);
	if (fu^fv) {
		ans+=2ll*siz[c][fu]*siz[c][fv];
		fa[c][fu]=fv;
		siz[c][fv]+=siz[c][fu];
	}
}

int main() {
	n=read();
	for (register int i=1;i<=n;++i) {
		fa[0][i]=fa[1][i]=i;
		siz[0][i]=siz[1][i]=1;
	}
	for (register int i=1;i<n;++i) {
		e[i].u=read();
		e[i].v=read();
		e[i].c=read();
		merge(e[i].u,e[i].v,e[i].c);
	}
	for (register int i=1;i<=n;++i) {
		int a=find(i,0),b=find(i,1);
//		ans+=1ll*siz[0][a]*siz[1][b]-1;
		ans+=1ll*(siz[0][a]-1)*(siz[1][b]-1);
	}
	cout<<ans<<endl;
	return 0;
}