CodeForces 1156 D.0-1-Tree(并查集)

D.0-1-Tree

You are given a tree (an undirected connected acyclic graph) consisting of n vertices and n−1 edges. A number is written on each edge, each number is either 0 (let’s call such edges 0-edges) or 1 (those are 1-edges).

Let’s call an ordered pair of vertices (x,y) (x≠y) valid if, while traversing the simple path from x to y, we never go through a 0-edge after going through a 1-edge. Your task is to calculate the number of valid pairs in the tree.

Input

The first line contains one integer n (2≤n≤200000) — the number of vertices in the tree.

Then n−1 lines follow, each denoting an edge of the tree. Each edge is represented by three integers xi, yi and ci (1≤xi,yi≤n, 0≤ci≤1, xi≠yi) — the vertices connected by this edge and the number written on it, respectively.

It is guaranteed that the given edges form a tree.

Output

Print one integer — the number of valid pairs of vertices.

Example

Input

7
2 1 1
3 2 0
4 2 1
5 2 0
6 7 1
7 2 1

Output

34

思路:

首先只走白色边和只走黑色边的点对合法
因此开两个并查集单独合并白边连通块和黑边连通块

对于每个同色的连通块
答案累加num*(num-1)

题目说走1边之后不能走0边，意思是可以走0边之后走1边，
对于每个点i，这该点属于的0连通块大小为num0，1连通块大小为1
则0的连通块走到1的连通块方案数为(num0-1)*(num1-1)，累加至答案
减1是因为要去掉自己

code：

#include <iostream>
#include <cstring>
#include <cstdio>
#include <cmath>
#include <algorithm>
using namespace std;
#define int long long
const int maxm=2e5+5;
int c[maxm],cc[maxm];//两种颜色的连通块大小
int f[maxm],ff[maxm];//两个并查集的pre数组
void init(){
    for(int i=0;i<maxm;i++){
        f[i]=ff[i]=i;
        c[i]=cc[i]=1;
    }
}
int ffind(int x,int *pre){
    return pre[x]==x?x:pre[x]=ffind(pre[x],pre);
}
void add(int a,int b,int ccc){//按颜色合并
    if(ccc==1){//黑色
        int x=ffind(a,f);
        int y=ffind(b,f);
        if(x!=y){
            f[x]=y;
            c[y]+=c[x];
        }
    }else{//白色
        int x=ffind(a,ff);
        int y=ffind(b,ff);
        if(x!=y){
            ff[x]=y;
            cc[y]+=cc[x];
        }
    }
}
signed main(){
    init();
    int n;
    cin>>n;
    for(int i=1;i<n;i++){
        int a,b,c;
        cin>>a>>b>>c;
        add(a,b,c);
    }
    int ans=0;
    for(int i=1;i<=n;i++){
        if(f[i]==i){//如果是黑根
            ans+=c[i]*(c[i]-1);//1->1
        }
        if(ff[i]==i){//如果是白根
            ans+=cc[i]*(cc[i]-1);//0->0
        }
        int x=ffind(i,f);//当前点属于的连通块的黑根
        int y=ffind(i,ff);//当前点属于的连通块的白根
        ans+=(c[x]-1)*(cc[y]-1);//0->1
    }
    cout<<ans<<endl;
    return 0;
}