人人都是好朋友(离散化 + 并查集)

离散化

#include<bits/stdc++.h>
using namespace std;
const int MAX = 2e6 + 5;
int n, f[MAX], a[MAX], b[MAX], c[MAX];
vector<int> d;
int find(int v) {
    return f[v] == v ? v : f[v] = find(f[v]);
}
void solve() {
    int t;
    scanf("%d", &t);
    int sums = 0;
    while (t--) {
        d.clear();
        scanf("%d", &n);
        sums += n;
        for (int i = 1; i <= 2 * n; ++i) f[i] = i;
        for (int i = 1; i <= n; ++i) {
            scanf("%d%d%d", &a[i], &b[i], &c[i]);
            d.push_back(a[i]);
            d.push_back(b[i]);
        }
        sort(d.begin(), d.end());//排序
        d.erase(unique(d.begin(), d.end()), d.end());//删除重复元素
        for (int i = 1; i <= n; ++i) {
            //索引元素离散化后对应的值
            a[i] = lower_bound(d.begin(), d.end(), a[i]) - d.begin() + 1;
            b[i] = lower_bound(d.begin(), d.end(), b[i]) - d.begin() + 1;
        }
        for (int i = 1; i <= n; ++i) {
            if (c[i] == 1) {
                int A = find(a[i]), B = find(b[i]);
                f[B] = A;
            }
        }
        bool flag = true;
        for (int i = 1; i <= n; ++i) {
            if (c[i] == 0) {
                int A = find(a[i]), B = find(b[i]);
                if (A == B) { flag = false; break; }
            }
        }
        printf("%s
", flag == false ? "NO" : "YES");
    }
}
int main(){
    solve();
    return 0;
}