How do I solve this using union-find algorithm? - c++

Question

Question: http://codeforces.com/contest/468/problem/B
Little X has n distinct integers: p1, p2, ..., pn. He wants to divide all of them into two sets A and B. The following two conditions must be satisfied:
If number x belongs to set A, then number a - x must also belong to set A.
If number x belongs to set B, then number b - x must also belong to set B.
Help Little X divide the numbers into two sets or determine that it's impossible.
Input
The first line contains three space-separated integers n, a, b (1 ≤ n ≤ 105; 1 ≤ a, b ≤ 109). The next line contains n space-separated distinct integers p1, p2, ..., pn (1 ≤ pi ≤ 109).
Output
If there is a way to divide the numbers into two sets, then print "YES" in the first line. Then print n integers: b1, b2, ..., bn (bi equals either 0, or 1), describing the division. If bi equals to 0, then pi belongs to set A, otherwise it belongs to set B.
If it's impossible, print "NO" (without the quotes).
Now, I developed the following code:
#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;
}
Basically, try to merge every element with either in Set A or in Set B. And then finally output NO if both of them are merged in turn. However, this gives a wrong answer on input as:
3 3 4
1 2 4
Output should be: NO whereas my code gives output as
YES
0 0 1
Where am I going wrong in my logic? Please help!

Answer 1

Consider the following cases for P being the set of all input numbers p[i]:
There is a number n1 in P that satisfies n1 = a - p[i] but no number n2 in P that satisfies n2 = b - p[i]
There is a number n1 in P that satisfies n1 = b - p[i] but no number n2 in P that satisfies n2 = a - p[i]
There is no number n in P that satisfies n = a - p[i] OR n = b - p[i]
There is a number n1 in P that satisfies n1 = a - p[i] AND a number n2 in P that satisfies n2 = b - p[i]
Whatever else may happen, if you run into situation 3. or 4. you want to report a failure ("NO").
If all numbers p[i] belong to case 1. or 2. the result should be valid.
You should check your code line by line, if it is compatible with the provided cases here.