0-1-Tree【Codeforces 1156D】【带权并查集】

题目链接

---

D. 0-1-Tree

time limit per test

2 seconds

memory limit per test

256 megabytes

input

standard input

output

standard output

You are given a tree (an undirected connected acyclic graph) consisting of ?n vertices and ?−1n−1 edges. A number is written on each edge, each number is either 00 (let's call such edges 00-edges) or 11 (those are 11-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 00-edge after going through a 11-edge. Your task is to calculate the number of valid pairs in the tree.

Input

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

Then ?−1n−1 lines follow, each denoting an edge of the tree. Each edge is represented by three integers ??xi, ??yi and ??ci (1≤??,??≤?1≤xi,yi≤n, 0≤??≤10≤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.

---

  题意：题意其实就是一句话，我们要先走0值的边，再走1值的边，问的是最多有几种走法？当然可以不走完的，走到哪里停下都是可以的。

  思路：想到并查集可能是在画图的过程中想到的，我们可以利用题中所给出的那张图来分析，是不是可以将所有的节点进行一个这样的处理，把每个节点的规划为由边权仅0和边权仅1到达的区分，并记录数量。我们分析从1号节点开始的节点的状态，发现1号节点出去只是链接着1权值的节点，记录其所有的1边权的点（包括自己本身）一共有5个，则从一号节点出发（1号节点只能作为出发的节点）的方案数则为（5-1=4）；再分析以2这号节点作为中转点的时候的情况。我们可以看到，2号节点直接链接的0权点的数目有2个，分别为3、5号节点，同时从2号节点直接可以通过1权边链接的点的数目为4个，分别为1、4、6、7，那么将2号节点作为中转站的时候，方案数为（3 * 5 - 1 = 14）表示着从0值的点出发到达1值的点的以及0到0和1到1的都要经过2的方案数。

  那么，我们只需要记录每个节点的权值为0所能到达的点以及权值为1所能到达的点的数量即可，但是0、1各自独立计算，这时候可以开一个二维的并查集"root[i][j]"表示"i为0或1"（边权）时候的"j节点的根节点的关系"，同时还可以记录下到根节点的所有的距离的和。

  最后的答案就是每个点的"siz[0][fid(i)] * siz[1][fid(i)] - 1;"之和。意义是从0权值节点出发经过i号点到达的全体1节点以及到达i节点的全体方案，所以最初的时候“siz[][] = 1;”；再则就是求从0号到达0号以及从1号到达1号的经过i 号节点的方案数。

---

#include <iostream>
#include <cstdio>
#include <cmath>
#include <string>
#include <cstring>
#include <algorithm>
#include <limits>
#include <vector>
#include <stack>
#include <queue>
#include <set>
#include <map>
#define lowbit(x) ( x&(-x) )
#define pi 3.141592653589793
#define e 2.718281828459045
#define INF 0x3f3f3f3f
#define HalF (l + r)>>1
#define lsn rt<<1
#define rsn rt<<1|1
#define Lson lsn, l, mid
#define Rson rsn, mid+1, r
#define QL Lson, ql, qr
#define QR Rson, ql, qr
#define myself rt, l, r
namespace fastIO {
#define BUF_SIZE 100000
    //fread -> read
    bool IOerror = 0;
    inline char nc() {
        static char buf[BUF_SIZE], *p1 = buf + BUF_SIZE, *pend = buf + BUF_SIZE;
        if(p1 == pend) {
            p1 = buf;
            pend = buf + fread(buf, 1, BUF_SIZE, stdin);
            if(pend == p1) {
                IOerror = 1;
                return -1;
            }
        }
        return *p1++;
    }
    inline bool blank(char ch) {
        return ch == ' ' || ch == '\n' || ch == '\r' || ch == '\t';
    }
    inline void read(int &x) {
        char ch;
        while(blank(ch = nc()));
        if(IOerror) return;
        for(x = ch - '0'; (ch = nc()) >= '0' && ch <= '9'; x = x * 10 + ch - '0');
    }
#undef BUF_SIZE
};
using namespace fastIO;
using namespace std;
typedef unsigned long long ull;
typedef long long ll;
const int maxN = 2e5 + 7;
int N, root[2][maxN], siz[2][maxN];
inline int fid(int *fa, int x) { return x == fa[x] ? x : fa[x] = fid(fa, fa[x]); }
inline void Merge(int *fa, int *siz, int x, int y)
{
    int u = fid(fa, x), v = fid(fa, y);
    if(u != v)
    {
        fa[u] = v;
        siz[v] += siz[u];
    }
}
inline void init()
{
    for(int i=1; i<=N; i++) root[0][i] = root[1][i] = i;
    for(int i=1; i<=N; i++) siz[0][i] = siz[1][i] = 1;
}
int main()
{
    scanf("%d", &N);
    //read(N);
    init();
    for(int i=1, u, v, w; i<N; i++)
    {
        scanf("%d%d%d", &u, &v, &w);
        //read(u);    read(v);    read(w);
        Merge(root[w], siz[w], u, v);
    }
    ll ans = 0;
    for(int i=1; i<=N; i++)
    {
        ans += (ll)siz[0][fid(root[0], i)] * (ll)siz[1][fid(root[1], i)] - 1;
    }
    printf("%lld\n", ans);
    return 0;
}