Educational Codeforces Round 64 (Rated for Div. 2)D(并查集，图)

#include<bits/stdc++.h>
using namespace std;
int f[2][200007],s[2][200007];//并查集，相邻点
int find_(int *f,int x){
    return f[x]==x?x:f[x]=find_(f,f[x]);
}
void add(int *f,int *s,int x,int y){
    x=find_(f,x);
    y=find_(f,y);
    if(x!=y)//不在一个联通块，将x并入y所在联通块中，将块内的相邻类型点数量更新
        f[x]=y,s[y]+=s[x];
}
int main(){
    int n;
    scanf("%d",&n);
    for(int i=1;i<=n;++i){
        f[0][i]=f[1][i]=i;
        s[0][i]=s[1][i]=1;
    }
    int x,y,z;
    for(int i=1;i<n;++i){
        scanf("%d%d%d",&x,&y,&z);
        add(f[z],s[z],x,y);
    }
    long long ans=0;
    for(int i=1;i<=n;++i)
        ans+=1ll*s[0][find_(f[0],i)]*s[1][find_(f[1],i)]-1;
        //从i点出发，终点为0或1结点，共有x-1+y-1，从0结点出发经过i终点为1结点，共有（x-1）*（y-1），合计x*y-1
    printf("%lld",ans);
    return 0;
}