c++ - 如何使用 union-find 算法解决这个问题？

Question

问题:http://codeforces.com/contest/468/problem/B

小 X 有 n 个不同的整数:p1,⟩p2,⟩...,⟩pn。他想把它们全部分成两组A和B，必须满足以下两个条件:

如果数字x属于集合A，那么数字a⟩-⟩x也一定属于集合A。 如果数字 x 属于集合 B，则数字 b⟩-⟩ x 也必须属于集合 B。 帮助小X将数字分成两组或确定不可能。

输入 第一行包含三个以空格分隔的整数 n, a,⟩b (1 ≤ n ≤⟩105; 1 ≤ a,⟩b⟩≤⟩109)。下一行包含 n 个以空格分隔的不同整数 p1, p2, ..., pn (1 ≤ pi ≤ 109)。

输出 如果有办法将数字分成两组，则在第一行打印“YES”。然后打印 n 个整数:b1,⟩b2,⟩...,⟩bn(bi 等于 0 或 1)，描述除法。如果bi等于0，则pi属于集合A，否则属于集合B。

如果不可能，打印“NO”(不带引号)。

现在，我开发了以下代码:

#include <iostream>
#include <cstdio>
#include <cstring>
#include <algorithm>
#include <set>
#include <map>

using namespace std;

int n,a,b,id;
int p[100100];
set<int> st;
map<int,int> mp;

int fa[100100];

int find(int x)
{
    if(x==fa[x]) return x;
    return fa[x]=find(fa[x]);
}

void Bing(int a,int b)
{
    int A=find(a),B=find(b);
    if(A==B) return ;
    fa[B]=A;
}

int main()
{
    scanf("%d%d%d",&n,&a,&b);
    for(int i=1;i<=n;i++)
    {
        scanf("%d",p+i);
        st.insert(p[i]);
        mp[p[i]]=++id;
        fa[i]=i;
    }
    fa[n+1]=n+1;///A
    fa[n+2]=n+2;///B

    for(int i=1;i<=n;i++)
    {
        int x=p[i];
        int flag = 0;
        if(st.count(a-x))
        {
            Bing(mp[x],mp[a-x]);
            flag = 1;
        }
        else
        {
            Bing(n+1,mp[x]);
            flag = 1;
        }
        if(st.count(b-x) && flag == 0)
        {
            Bing(mp[x],mp[b-x]);
        }
        else if (flag == 0)
        {
            Bing(n+2,mp[x]);
        }
    }

    if(find(n+1)==find(n+2))
    {
        puts("NO");
    }
    else
    {
        puts("YES");
        for(int i=1;i<=n;i++)
        {
            printf("%d ",find(i)==find(n+1));
        }
        putchar(10);
    }

    return 0;
}

基本上，尝试将每个元素与 Set A 或 Set B 中的一个元素合并。然后如果它们都依次合并，则最终输出 NO。但是，这会给出错误的输入答案:

3 3 4
1 2 4

输出应为:NO 而我的代码给出的输出为

YES
0 0 1 

我的逻辑哪里出错了？请帮忙!

Answer 1

考虑以下情况，P 是所有输入数字 p[i] 的集合:

1. P 中有一个数 n1 满足 n1 = a - p[i] 但 P 中没有数 n2 满足 n2 = b - p[i]
2. P 中有一个数 n1 满足 n1 = b - p[i] 但 P 中没有数 n2 满足 n2 = a - p[i]
3. 没有满足 n = a - p[i] OR n = b - p[i] 的数 n in P
4. P 中有一个数 n1 满足 n1 = a - p[i] 和一个数 n2 in P 满足 n2 = b - p[i]

无论发生什么情况，如果您遇到情况 3. 或 4. 您想要报告失败(“否”)。

如果所有数字 p[i] 都属于情况 1. 或 2. 结果应该是有效的。

您应该逐行检查您的代码，是否与此处提供的案例兼容。