牛客小白月赛24 H.人人都是好朋友(离散化&并查集）

题目传送门

思路：将朋友的关系建立一个并查集。再遍历一遍，看两个敌人的根结点是否相同。

AC代码：

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int N=1e6+5;
#define mst(a) memset(a,0,sizeof a)
int s[N<<1],a[N],b[N],d[N<<1],h[N<<1];
int c[N];
int find(int x){
	if(x!=s[x]) s[x]=find(s[x]);
	return s[x];
}
void union_set(int x,int y){
	 x=find(x),y=find(y);
	 if(h[x]==h[y]){
	 	h[y]++;
	 	s[x]=y;
	 }
	 else  h[x]>h[y]?s[y]=x:s[x]=y;
}
int read(){ //快读 
	int s=0,w=1;
	char ch=getchar();
	while(ch<'0'||ch>'9') {if(ch=='-') w=-1;ch=getchar();}
	while(ch>='0'&&ch<='9') s=s*10+ch-'0',ch=getchar();
	return s*w;
}

int main(){
	int t;
	t=read();
	while(t--){
		int n,k=0;
		n=read();
		for(int i=1;i<=n;i++)
		{
			a[i]=read(),b[i]=read(),c[i]=read();
			d[++k]=a[i],d[++k]=b[i];
			s[i]=i,s[i+n]=i+n,h[i]=h[i+n]=0;
		}
		sort(d+1,d+k+1);//排序后离散化 
		int cnt=unique(d+1,d+k+1)-d;
		for(int i=1;i<=n;i++) 
		{
			a[i]=lower_bound(d+1,d+cnt+1,a[i])-d;
			b[i]=lower_bound(d+1,d+cnt+1,b[i])-d;	
		}
		for(int i=1;i<=n;i++) 
			if(c[i])	union_set(a[i],b[i]); //将所有友好关系建立一个并查集. 
		int f=1; 
		for(int i=1;i<=n;i++)
			if(!c[i]&&s[a[i]]==s[b[i]]) //找矛盾的即可 
			{
				f=0;
				break;
			}
		puts(f?"YES":"NO");
	}
	return 0;
}