清华大学算法训练营——程序六：等式

等式（并查集）

时间限制：2s，空间256MB

问题描述

n个变量和m个“相等”或“不相等”的约束条件，请你判定是否存在一种赋值方案满足所有m个约束条件。

输入

第一行一个整数T，表示数据组数。(T<=100)

接下来会有T组数据，对于每组数据：

第一行是两个整数n,m，表示变量个数和约束条件的个数。(1<=n,m<=500000)

接下来m行，每行三个整数a,b,e(1<=a,b<=n)，表示第a个变量和第b个变量的关系：

若e=0则表示第a个变量不等于第b个变量；
若e=1则表示第a个变量等于第b个变量

保证所有数据的n总和不超过500000。

输出

输出T行，第i行表示第i组数据的答案。若第i组数据存在一种方案则输出”Yes”；否则输出”No”（不包括引号）。

输入样例

输出样例

Yes
No

样例解释

对于第一组数据，有5个约束：

变量1=变量2
变量2=变量3
变量3=变量4
变量1=变量4
变量2≠变量5

显然我们可以令：

变量1=变量2=变量3=变量4=任意一个数值
变量5=任意一个和变量2不同的数值

故第一组数据输出”Yes”。

对于第二组数据，有3个约束：

变量1=变量2
变量2=变量3
变量1≠变量3

由前两个约束可推出变量1=变量3，但第三个约束表明变量1≠变量3，矛盾。

故第二组数据输出”No”。

#include <bits/stdc++.h>
using namespace std;
// ================= 代码实现开始 =================
const int N = 300005;
// Fanther：每个节点的父亲节点
// Rank:节点的秩（用于启发式合并）
int Father[N],Rank[N];

//查找节点x所在集合的根
int find(int x){
    return Father[x]==x?x:Father[x]=find(Father[x]);
}
// n：变量个数
// m：约束个数
// A：大小为m的数组，表示m条约束中的a
// B：大小为m的数组，表示m条约束中的b
// E：大小为m的数组，表示m条约束中的e
// 返回值：若能找出一种方案，返回"Yes"；否则返回"No"
string getAnswer(int n, int m, vector<int> A, vector<int> B, vector<int> E) {
    for(int i=1; i<=n; ++i){
        Father[i] = i;
        Rank[i] = 0;
    }

     int cnt = 0;
     for(int i=0; i<m; ++i)  //将e=1的操作提到e=0的操作前
        if(E[i] == 1) {
           swap(E[i], E[cnt]);
           swap(A[i], A[cnt]);
           swap(B[i], B[cnt]);
           ++cnt;
     }
     for(int i=0; i<m; ++i){
        int setA = find(A[i]);
        int setB = find(B[i]);
        if(E[i] == 0){
            if(setA == setB){
                return "No";
            }
        }
        // E[i]==1 即 A[i]与B[i]有共同的父亲节点
        else{
            //在Union时把Rank更高的父亲节点作为根节点
            if(setA != setB){
                if(Rank[setA] > Rank[setB])
                    swap(setA,setB);
                // Rank[setB]的等级更高
                Father[setA] = setB;
                if(Rank[setA] == Rank[setB])
                    Rank[setB] +=1;
            }
        }
     }
     return "Yes";
}

// ================= 代码实现结束 =================

int main() {
    int T;
    for (scanf("%d", &T); T--; ) {
        int n, m;
        scanf("%d%d", &n, &m);
        vector<int> A, B, E;
        for (int i = 0; i < m; ++i) {
            int a, b, e;
            scanf("%d%d%d", &a, &b, &e);
            A.push_back(a);
            B.push_back(b);
            E.push_back(e);
        }
        printf("%s\n", getAnswer(n, m, A, B, E).c_str());
    }
    return 0;
}