题解 | 人人都是好朋友

#include <cassert>
#include <iostream>
#include <unordered_map>
#include <vector>
#include <bits/stdc++.h>
using namespace std;
struct DSU {
    vector<int> fa, sz;
    int component;

    DSU(int n, int type = 0) : fa(n + type), sz(n + type, 1) {
        iota(fa.begin(), fa.end(), 0); // fa[i] = i
        component = n;
    }

    int find(int x) {
        return fa[x] == x ? x : fa[x] = find(fa[x]);
    }

    void merge(int a, int b) {
        a = find(a), b = find(b);
        if (a == b)
            return;
        if (sz[a] < sz[b])
            swap(a, b);
        fa[b] = a;
        sz[a] += sz[b];
        component--;
    }
};

bool solve(int n) {
    unordered_map<int,int> id;
    int cnt = 0;
    DSU dsu1(2*n);
    for(int i=0;i<n;i++){
        int a,b,c;
        cin >> a >> b >> c;
        if(!id.count(a)) id[a]=cnt++;
        if(!id.count(b)) id[b]=cnt++;
        a=id[a];
        b=id[b];
        // assert(a<n&&b<n);
        if(c==1)dsu1.merge(a,b);
        if(c==0&&dsu1.find(a)==dsu1.find(b))return false;
    }
    return true;
}

int main() {
    ios::sync_with_stdio(false),cin.tie(nullptr);
    int t;
    cin >> t;
    while (t--) {
        int n;
        cin >> n;
        cout << (solve(n) ? "YES" : "NO")<<endl;
    }
}
// 64 位输出请用 printf("%lld")

离散化+并查集