CF1156D 0-1-Tree

路径考虑顺序。

显然合法的路径只有以下两种：

1. 一段 (0) 加一段 (1) 或一段 (1) 加一段 (0)。
2. 全 (0) 或全 (1)。

用并查集将边权为 (0) 和 (1) 的边分别缩起来，对于一个大小为 (siz) 的联通块，第二种的答案是 (siz(siz-1))。

对于第一种，我们枚举 (0) 和 (1) 的分界点，它所在 (0) 和 (1) 的联通块大小分别是 (siz_0) 和 (siz_1)，它的答案是 ((siz_0-1)(siz_1-1))。

然后这题就做完了，时间复杂度 (O(nalpha(n)))。

code:

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
#define N 2
#define NN 200005
#define For(i,x,y)for(i=x;i<=(y);i++)
int fa[NN][N],siz[NN][N];
int read()
{
	int A;
	bool K;
	char C;
	C=A=K=0;
	while(C<'0'||C>'9')K|=C=='-',C=getchar();
	while(C>'/'&&C<':')A=(A<<3)+(A<<1)+(C^48),C=getchar();
	return(K?-A:A);
}
int find(int p,bool type)
{
	if(fa[p][type]!=p)fa[p][type]=find(fa[p][type],type);
	return fa[p][type];
}
inline void unite(int p,int q,bool type)
{
	p=find(p,type),q=find(q,type);
	if(p==q)return;
	if(siz[p][type]<siz[q][type])fa[p][type]=q,siz[q][type]+=siz[p][type];
	else fa[q][type]=p,siz[p][type]+=siz[q][type];
}
int main()
{
	bool c;
	ll ans=0;
	int n,i,x,y;
	n=read();
	For(i,1,n)fa[i][0]=fa[i][1]=i,siz[i][0]=siz[i][1]=1;
	For(i,1,n-1)
	{
		x=read(),y=read(),c=read();
		unite(x,y,c);
	}
	For(i,1,n)ans+=1LL*(siz[find(i,0)][0]-1)*(siz[find(i,1)][1]-1);
	For(i,1,n)
	{
		ans+=1LL*siz[find(i,0)][0]*(siz[find(i,0)][0]-1)+1LL*siz[find(i,1)][1]*(siz[find(i,1)][1]-1);
		siz[find(i,0)][0]=siz[find(i,1)][1]=0;
	}
	cout<<ans;
	return 0;
}