集合问题（并查集）

链接：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

#include <iostream>
#include <cstring>
#include <string>
#include <algorithm>
#include <cstdio>
#include <map>
using namespace std;

const int maxn = 1e5 + 7;
int n, A, B, fa[maxn], a[maxn];
map<int, int>mp;
void init() {
    for (int i = 0; i <= n + 1; i++)
        fa[i] = i;
}
int find(int x)
{
    if (fa[x] == x)return x;
    else return fa[x] = find(fa[x]);
}
void uni(int x, int y)
{
    int px = find(x);
    int py = find(y);
    if (px == py)return;
    else fa[px] = py;
}
int main() 
{
    scanf("%d%d%d", &n, &A, &B);
    int max1 = 0;
    for (int i = 1; i <= n; i++)
    {
        scanf("%d", &a[i]);
        mp[a[i]] = i;
        max1 = max(max1, a[i]);
    }
    if (max1 >= max(A, B))
    {
        return printf("NO
"), 0;
    }
    else
    {
        init();
        for (int i = 1; i <= n; i++)
        {
            if (mp[B - a[i]])//如果B-a[i]存在
                uni(i, mp[B - a[i]]);//就把（B-a[i]）数的序号和i合并在一起
            else uni(i, 0);//否则放到和0一个集合（a）里去
            if (mp[A - a[i]])
                uni(i, mp[A - a[i]]);
            else uni(i, n + 1);
        }
        int af = find(0);
        int bf = find(n + 1);
        if (af == bf)//一个数两个集合里都必须有，产生矛盾
        {
            return printf("NO
"), 0;
        }
        else
        {
            printf("YES
");
            for (int i = 1; i <= n; i++)
            {
                if (i != 1)printf(" ");
                if (af == find(i))printf("0");
                else printf("1");
            }
            printf("
");
        }
    }
    return 0;
}