codeforces1156D 0-1-Tree 并查集

题目传送门

题意：

　　给定一棵n个点的边权为0或1的树，一条合法的路径(x,y)(x≠y)满足，从x走到y，一旦经过边权为1的边，就不能再经过边权为0的边，求有多少边满足条件？

思路：

　　首先这道题，换根dp也可以过（树形dp，点这里）

　　那么如何并查集做呢，我们考虑一个点$u$，我们将与u通过0边相连的连通点的数量记做$siz0$，将与u通过1边相连的连通点的数量记做$siz1$，那么所有符合条件的0-1边将u作为转折点的，0-0边将u作为终点的，1-1边将u作为终点的边的数量就等于$siz0*siz1-1$，减去的1就是自己连到自己。

　　用并查集维护上面的信息，即$f[u][0]$表示u的0边祖先，$f[u][0]$表示u的1边祖先，然后$merge$即可。

#pragma GCC optimize (2)
#pragma G++ optimize (2)
#pragma comment(linker, "/STACK:102400000,102400000")
#include<bits/stdc++.h>
#include<cstdio>
#include<vector>
#define rep(i,a,b) for(int i=a;i<=b;i++)
#define dep(i,b,a) for(int i=b;i>=a;i--)
#define clr(a,b) memset(a,b,sizeof(a))
#define pb push_back
#define pii pair<int,int >
using namespace std;
typedef long long ll;
const int maxn=;
ll rd()
{
    ll x=,f=;char ch=getchar();
    while(ch<''||ch>''){if(ch=='-')f=-;ch=getchar();}
    while(ch>=''&&ch<=''){x=x*+ch-'';ch=getchar();}
    return x*f;
}
int f[maxn][],siz[maxn][];
int n,a,b,c;
int find(int x,int y){
    if(!f[x][y]){
        return x;
    }
    return f[x][y]=find(f[x][y],y);
}
int main(){
    cin>>n;
    rep(i,,n-){
        scanf("%d%d%d",&a,&b,&c);
        int fx=find(a,c);
        int fy=find(b,c);
        if(fx!=fy){
            f[fx][c]=fy;
        }
    }
    rep(i,,n){
        siz[find(i,)][]++,siz[find(i,)][]++;
    }
    ll ans=;
    rep(i,,n){
        ans+=1ll*siz[find(i,)][]*siz[find(i,)][]-;
    }
    cout<<ans<<endl;
}