20-集合问题(并查集)

链接：https://www.nowcoder.com/acm/contest/77/D
来源：牛客网

时间限制：C/C++ 1秒，其他语言2秒
空间限制：C/C++ 262144K，其他语言524288K
64bit IO Format: %lld

题目描述

给你a,b和n个数p[i]，问你如何分配这n个数给A,B集合，并且满足：

若x在集合A中，则a-x必须也在集合A中。

若x在集合B中，则b-x必须也在集合B中。

输入描述:

第一行 三个数 n a b  1<=n<=1e5  1<=a,b<=1e9
第二行 n个数 p1 p2 p3...pn 1<=pi<=1e9

输出描述:

如果可以恰好分开就输出第一行 YES
然后第二行输出 n个数 分别代表pi 是哪个集合的  0 代表是A集合 1代表是B 集合
不行就输出NO
放在哪个集合都可以的时候优先放B

示例1

输入

4 5 9
2 3 4 5

输出

YES
0 0 1 1

示例2

输入

3 3 4
1 2 4

输出

NO

题意：

给你a,b和n个数p[i]，问你如何分配这n个数给A,B集合，并且满足：

若x在集合A中，则a-x必须也在集合A中。

若x在集合B中，则b-x必须也在集合B中。

思路：

对于一个x来说，能分成以下4种情况（这里a!=b）：

1.a-x不存在，b-x不存在。这种情况直接输出NO。

2.a-x存在，b-x不存在。这种情况把x和a-x放在集合A中。

3.a-x不存在，b-x存在。这种情况把x和b-x放在集合B中。

4.a-x存在，b-x存在。这种情况比较我们就不能直接确定要放A还是要放B了。

那么我们接下还需要判断是否存在y有a-y=b-x或者b-y=a-x其中之一存在。为什么是其中之一呢？

假设有a-y=b-x和b-y=a-x同时存在，那么将两式子相减有x-y=y-x -> x*2=y*2 -> x=y成立，则有a-y=a-x=b-x -> a=b矛盾。所以上面两种情况只可能存在一种或都不存在。如果存在a-y=b-x，那么就把x和a-x和y和a-y(b-x)都放到集合A中。同理b-y=a-x存在的情况。

如果都不存在，那么我们就无法决则把x和a-x和b-x放到集合A中还是集合B中，因为放到那一边都会剩下另一个无法去放，所以就输出NO。这里用并查集来维护即可。

#include <bits/stdc++.h>
using namespace std;
const int Max = 1e5 + 5;
map <int, int> mp;
int fa[Max], f[Max], n, a, b;

void init(){
	for(int i = 0; i <= n + 1; i++){  //注意初始化的范围0~n+1 
		fa[i] = i;
	}
}
/*
int find(int x){    
	int i = x;
	while(i != fa[i]){
		i = fa[i];
	}
	int j = x;
	while(fa[j] != i){     //数据压缩 
		int k = fa[j];
		fa[j] = i;
		j = k;
	}
	return i;
}
*/
int find(int x){  
	if(fa[x] == x)
		return x;
	else 
		return fa[x] = find(fa[x]);   //也压缩了，更简洁 
} 

void Union(int x, int y){
	int i = find(x);
	int j = find(y);
	if(i != j){
		fa[i] = j;
	}
}

int main(){
	cin >> n >> a >>b;
	int mymax = -1;
	for(int i = 1; i <= n; i++){
		cin >> f[i];
		mp[f[i]] = i;
		mymax = max(mymax, f[i]);
	}
	if(mymax > max(a, b)){
		cout << "NO" << endl;
	}
	else{
		init();
		for(int i = 1; i <= n; i++){
			if(mp[b - f[i]])			 //如果b-f[i]存在	
				Union(i, mp[b - f[i]]);  //b-f[i]存在就i和b-f[i]的映射放在一个集合里面 
			else						
				Union(i, 0);			 //不存在是就只能讲i放入到0集合即(a集合)
			if(mp[a - f[i]])
				Union(i, mp[a - f[i]]);
			else
				Union(i, n + 1); 
		}
		int f1 = find(0);
		int f2 = find(n + 1);
		if(f1 == f2){			//一个数两个集合里都必须有，产生矛盾
			cout << "NO" << endl;
		}
		else{
			cout << "YES" << endl;
			for(int i = 1; i <= n; i++){
				if(i != 1)
					cout << " ";
				if(find(i) == f1){    //f[i]的映射和0在同一个集合，这输出0，表示在A集合 
					cout << 0;
				}
				else
					cout << 1;
			}
			cout << endl;
		}
	}
	return 0;
} 
/*
可以手动走一遍： 
6 5 9
2 3 4 5 1 8
*/