哥大周赛题目 0-1 Tree （BFS + 并查集）

上周本地比赛，老wf选手都退役了，只剩我们这些22届本科升研究生来参赛了

题目不是很难，11题，比之前的训练赛要简单很多，开场A了4题签到 + 1裸dp + 1做过 + 1xjb乱搞。结果最后一题本来没思路，但是出去上了个厕所之后回来突发奇想，用排序大法乱搞了一波，wa9，然后我接着又写了三个排序，每个都试一遍，结果莫名其妙冲过去了，最后8题结束。

然后意外的发现今年还能打，走向了未曾设想的道路，即代表学校出战，研一还继续打明年三月的regional，神奇的体验

拐回来看题，首先今天有点事，做完溜了

这题给出一棵带权无根树，每个边要么是0要么是1，问图中存在多少对点(u, v)，使得u到v的路径上不存在先走1再走0, 另外(u, v)和(v, u)算两对

首先题意很清楚，我们很容易想到，全0和全1都是合法的，先走0再走1也是合法的

那么首先我们把三种情况列出来

然后我们再把所有全0和全1的连通分量处理出来，用并查集记录其大小，因为树的性质，所以每个子集对答案的贡献就是(size * (size - 1))

然后我们再从每个点开始bfs，看看0和哪些1边接壤了，有了上一步的处理，我们可以O(1)时间内得知相邻的点集的大小

所以每次对答案的贡献就是(size0 - 1) * (size1 - 1)， -1是因为接壤的点被重复计算了，所以必须要减掉

然后注意在ll乘法的时候作为右值一定要*1ll，不然会挂

#include <bits/stdc++.h>
using namespace std;
constexpr int limit =  (3000000 + 5);//防止溢出
#define INF 0x3f3f3f3f
#define inf 0x3f3f3f3f3f
#define lowbit(i) i&(-i)//一步两步
#define EPS 1e-9
#define FASTIO  ios::sync_with_stdio(false);cin.tie(0),cout.tie(0);
#define ff(a) printf("%d\n",a );
#define pi(a, b) pair<a,b>
#define rep(i, a, b) for(ll i = a; i <= b ; ++i)
#define per(i, a, b) for(ll i = b ; i >= a  ; --i)
#define MOD 998244353
#define traverse(u) for(int i = head[u]; ~i ; i = edge[i].next)
#define FOPEN freopen("C:\\Users\\tiany\\CLionProjects\\akioi\\data.txt", "rt", stdin)
#define FOUT freopen("C:\\Users\\tiany\\CLionProjects\\akioi\\dabiao.txt", "wt", stdout)
typedef long long ll;
typedef unsigned long long ull;
char buf[1 << 23], *p1 = buf, *p2 = buf, obuf[1 << 23], *O = obuf;

inline ll read() {
#define getchar() (p1==p2&&(p2=(p1=buf)+fread(buf,1,1<<21,stdin),p1==p2)?EOF:*p1++)
    ll sign = 1, x = 0;
    char s = getchar();
    while (s > '9' || s < '0') {
        if (s == '-')sign = -1;
        s = getchar();
    }
    while (s >= '0' && s <= '9') {
        ll  x = (x << 3) + (x << 1) + s - '0';
        s = getchar();
    }
    return x * sign;
#undef getchar
}//快读
void print(ll x) {
    if (x / 10) print(x / 10);
    *O++ = x % 10 + '0';
}

void write(ll x, char c = 't') {
    if (x < 0)putchar('-'), x = -x;
    print(x);
    if (!isalpha(c))*O++ = c;
    fwrite(obuf, O - obuf, 1, stdout);
    O = obuf;
}

int n, m, k;
int a[limit];
int vis[limit];
int fa[limit];
int fa2[limit];
int get_fa(int x){
    return x == fa[x] ? x : fa[x] = get_fa(fa[x]);
}
int get_fa2(int x){
    return x == fa2[x] ? x : fa2[x] = get_fa2(fa2[x]);
}
void init(){
    rep(i,1,n){
        fa[i] = i;
        fa2[i] = i;
    }
}
void merge(int x, int y){
    if(x > y)swap(x,y);
    int pa = get_fa(x);
    int pb = get_fa(y);
    fa[pb] = pa;
}
void merge0(int x, int y){
    if(x > y)swap(x,y);
    int pa = get_fa2(x);
    int pb = get_fa2(y);
    fa2[pb] = pa;
}
vector<pi(int, int)>g[limit];
void solve(){
    cin>>n;
    init();
    rep(i,1,n - 1){
        int x,y,w;
        cin>>x>>y>>w;
        g[x].push_back({y,w});
        g[y].push_back({x,w});
    }
    ll ans = 0;
    auto bfs = [&](int x, int val = 1) -> ll {
        queue<int>q;
        int res = 0;
        vis[x] = 1;
        q.push(x);
        while(q.size()){
            int u = q.front();
            if(val)merge(x, u);
            else merge0(x, u);
            q.pop();
            ++res;
            for(auto & [v, w]: g[u]){
                if(w == val and !vis[v]){
                    vis[v] = 1;
                    q.push(v);
                }
            }
        }
        return res * (res -  1);
    };
    map<int, int>one, zero;
    auto bfs2 = [&](int x) -> ll {
        set<int>s;
        queue<int>q;
        ll res = 0;
        vis[x] = 1;
        q.push(x);
        int flag = 0;
        int flag2 = 0;
        while(q.size()){
            int u = q.front();
            q.pop();
            for(auto & [v, w]: g[u]){
                if(!w and !vis[v]){
                    vis[v] = 1;
                    q.push(v);
                    flag = 1;
                }else{
                    s.insert(fa[v]);
                    flag2 = 0;
                }
            }
        }
        if(flag){
            for(auto & it : s){
                res += (1ll * zero[get_fa2(x)] - 1) * (one[get_fa(it)] - 1);
            }
        }
        return res;
    };
    rep(i,1,n){
        if(!vis[i]){
            bfs(i);
        }
    }
    fill(vis, vis + 1 + n, 0);
    rep(i,1,n){
        if(!vis[i]){
            bfs(i, 0);
        }
    }

    rep(i,1,n){
        one[get_fa(i)]++;
        zero[get_fa2(i)]++;
    }
    for(const auto & [key, val] : one){
        ans += (1ll * val * (val - 1));
    }
    for(const auto & [key, val] : zero){
        ans += (1ll * val * (val - 1));
    }
    fill(vis, vis + 1 + n, 0);
    rep(i,1,n){
        if(!vis[i]){
            ll res = bfs2(i);
            ans += 1ll * res;
        }
    }
    cout<<ans<<endl;

};
int32_t main() {
#ifdef LOCAL
    FOPEN;
//    FOUT;
#endif
    FASTIO
//    int kase;
//    cin>>kase;
//    while (kase--)
    solve();
    cerr << "Time elapsed: " << 1.0 * clock() / CLOCKS_PER_SEC << "s\n";
    return 0;
}