牛客-集合问题——并查集的运用

题目

题目链接

题目解析

简述：将下标（用哈希表根据值来维护下标信息）用并查集连成集合，设置两个单独结点a,b来进行额外的维护。一旦无法连入A则把该集合整个根结点设为b，一旦无法连入B则把该集合整个根节点设为a。

解题代码

#include <bits/stdc++.h>

using namespace std;
#define MAXN 100010
int n, A, B, a, b, maxn = 0, val[MAXN], fa[MAXN];
unordered_map<int, int> pos;//TODO 记录值对应的下标

//TODO 并查集查找操作
int find(int x) { return fa[x] == x ? x : fa[x] = find(fa[x]); }

//TODO 并查集合并操作
void merge(int x, int y) {
    x = find(x);
    y = find(y);
    if (x != y)fa[y] = x;
}

void work() {
    if (maxn >= max(A, B)) {//TODO 如果最大的值比A/B都大则不可能放入A/B集合
        printf("NO\n");
        return;
    }
    for (int i = 1; i <= n; i++) {//TODO 由于此题给出的数据应该是不存在重复元素。否则需要visit对下标进行优化
        if (pos[B - val[i]])
            merge(i, pos[B - val[i]]);
        else merge(a, i);
        if (pos[A - val[i]])
            merge(i, pos[A - val[i]]);
        else merge(b, i);
    }
    a = find(a);
    b = find(b);//TODO 如果a，b的根结点相同，则存在一个元素在两个集合里，或者一个元素不在任何集合里
    if (a == b)printf("NO\n");
    else {
        printf("YES\n");
        for (int i = 1; i <= n; i++) {
            //TODO 根据它下标所在的集合来打印
            if (find(i) == a)printf("0 ");
            else printf("1 ");
        }
        printf("\n");
    }

}

void init() {
    ios::sync_with_stdio(false);
    cin >> n >> A >> B;
    a = n + 1;
    b = n + 2;//TODO 创建两个与普通下标有别的位置，表示两个集合的根点
    fa[a] = a;
    fa[b] = b;
    for (int i = 1; i <= n; i++) {
        cin >> val[i];
        fa[i] = i;
        pos[val[i]] = i;
        maxn = max(maxn, val[i]);
    }
}

int main() {
    init();
    work();
    return 0;
}