题解 | 人人都是好朋友

#include<bits/stdc++.h>
#define endl '\n'
using namespace std;

class dsu {
    vector<int> par;
    vector<int> rank;
  public:
    dsu(int n): par(n + 1), rank(n + 1) {
        for (int i = 0; i <= n; i++) {
            par[i] = i;
            rank[i] = 1;
        }
    }

    int find(int k) {
        if (par[k] == k) return k;
        return par[k] = find(par[k]);
    }

    void unit(int x, int y) {
        x = find(x);
        y = find(y);
        if (x == y) return;
        if (rank[x] < rank[y]) swap(x, y);
        par[y] = x;
        rank[x] += rank[y];
    }

    inline bool same(int x, int y) {
        return find(x) == find(y);
    }
};

void solve() {
    int n;
    cin >> n;
    vector<int> ai(n);
    vector<int> bi(n);
    vector<int> ci(n);
    vector<int> d(2*n);
    int cnt = -1;
    for (int i = 0; i < n; i++) {
        cin >> ai[i] >> bi[i] >> ci[i];
        d[++cnt] = ai[i];
        d[++cnt] = bi[i];
    }
    sort(d.begin(),d.end());
    auto ed = unique(d.begin(), d.end());
    for (int i = 0; i < n; i++)
    {
        ai[i] = lower_bound(d.begin(),ed,ai[i]) - d.begin();
        bi[i] = lower_bound(d.begin(),ed,bi[i]) - d.begin();
    }
    dsu tree1(ed-d.begin());
    dsu tree2(ed-d.begin());
    for(int i = 0;i<n;i++) {
        if(ci[i] == 1) {
            if(tree2.same(ai[i], bi[i])) return void(cout << "NO" << endl);
            tree1.unit(ai[i],bi[i]);
        } else {
            if(tree1.same(ai[i], bi[i])) return void(cout << "NO" << endl);
            tree2.unit(ai[i],bi[i]);
        }
    }
    return void(cout << "YES" << endl);
}

int main() {
    cin.tie(nullptr)->sync_with_stdio(false);

    int t;
    cin >> t;
    while (t--) {
        solve();
    }
    return 0;
}

并查集加上离散化，注意离散化的去重操作