Educational Codeforces Round 64 (Rated for Div. 2)D. 0-1-Tree（思维并查集）

link

题意

定义（x，y）表示x到y，其中我们经过1边后不能经过0边，求有多少对（x，y）

我们维护两个联通块 0/1
考虑两种情况

 如果只有0 / 1 它的贡献  siz[0][i]*(siz[0][i]-1)
另一种情况是先0后1  /先1后0    贡献为 (siz[0][x]-1)*(siz[1][y]-1)
用并查集维护就行

#include <bits/stdc++.h>
using namespace  std;
#define  int long long
//typedef long long ll;
typedef pair<int,int> pii;
#define x first 
#define y second
#define pb  push_back
#define inf 1e18
#define IOS   std::ios::sync_with_stdio(false);cin.tie(0);cout.tie(0);
#define  fer(i,a,b)  for(int i=a;i<=b;i++)
#define  der(i,a,b)  for(int i=a;i>=b;i--) 
 
const int mod=1e9+7;
/*ll qsm(int a,int b)
{ll res=1;  while(b){  if(b&1) res=res*a%mod;  a=a*a%mod;  b>>=1; } return res;}
*/
const int N=4e5+10;
int n,m;
int k;
int a[N];
int f[2][N];
int siz[2][N];
int find(int id,int x)
{
	return f[id][x]==x?x:f[id][x]=find(id,f[id][x]);
}
void solve()
{
	cin>>n;
    fer(i,1,n)
    {
   	f[0][i]=i;f[1][i]=i;
   	siz[0][i]=siz[1][i]=1;
    }
   fer(i,1,n-1)
   {
   	int a,b,c;
   	cin>>a>>b>>c;
   	int x=find(c,a);
   	int y=find(c,b);
   	if(x==y)continue;
   	f[c][x]=y;
   	siz[c][y]+=siz[c][x]; 	
   }
   int ans=0;
   fer(i,1,n)
   {
   	if(f[0][i]==i) ans+=siz[0][i]*(siz[0][i]-1);
   	if(f[1][i]==i)  ans+=siz[1][i]*(siz[1][i]-1);
   	int x=find(0,i);
   	int y=find(1,i);
   	ans+=(siz[0][x]-1)*(siz[1][y]-1);
   }
   cout<<ans<<endl;
}

signed main()
{
	IOS;
	int _=1;
    
	//cin>>_;
	while(_--) solve();
	return 0;
}