C++ 代码

#include <iostream>
#include <vector>
#include <unordered_map>
using namespace std;
using PII = pair<int, int>;

const int N = 2e6 + 10;
int p[N];
int cnt;
unordered_map<int, int> m;

int get(int x){ // 离散化， 将 10^9 映射到 2 * 10^6
    if(m.count(x) == 0) m[x] = ++ cnt;
    return m[x];
}

int find(int a){
    if(p[a] != a) p[a] = find(p[a]);
    return p[a];
}

void union1(int a, int b)
{
    int pa = find(a), pb = find(b);
    if(pa != pb) p[pa] = pb;
}

int main(){
    int t;
    cin >> t;
    while(t -- ){
        cnt = 0;
        m.clear();
        int n;
        cin >> n;
        vector<PII> v1, v2;
        for(int i = 0; i < n; ++ i){
            int x, y, e;
            scanf("%d %d %d", &x, &y, &e);
            if(e == 1) v1.emplace_back(get(x), get(y));
            else v2.emplace_back(get(x), get(y));
        }

        for(int i = 1; i <= cnt; ++ i) // 并查集初始化
            p[i] = i;

        for(auto & v: v1) // 合并相等的变量
            union1(v.first, v.second);

        bool ret = true;
        for(auto & v: v2) // 依次判断是否出现矛盾
            if(find(v.first) == find(v.second)){
                ret = false;
                break;
            }

        if(ret) puts("YES");
        else puts("NO");
    }
    return 0;
}