Codeforces 469D Two Sets

题目链接

https://codeforces.com/contest/469/problem/D

题目大意

给你一个长度为 N 的序列 p 和两个数 a , b 

现有两个集合 A , B , 要求你将这个序列放入集合A 、B（每个数只能放入一个集合）使得

如果 pi 存在于 A ， 则 a - pi 也存在于 A ， 如果 pi 存在于 B ，则 b - pi也存在于 B

解题思路

思维 + 并查集

我们定义 rootA 代表集合 A , rootB 代表集合 B

对于 pi ，我们分四种情况讨论

①、如果 a - pi 存在 ，则 pi 和 a - pi 必将存在于相同集合

②、如果 a - pi 不存在，则 pi 必将存在于集合 B , 即 pi 与  rootB 存在于相同集合

③、如果 b - pi 存在 ，则 pi 和 b - pi必将存在于相同集合

④、如果 b - pi 不存在，则 pi 必将存在于集合 A , 即 pi 与 rootA 存在于相同集合

以上情况我们用并查集维护 pi 下标来操作，其中 rootA 我们定义为 n + 1 , rootB定义为 n + 2

那么最后如果 find（rootA）== find（rootB）则答案不存在 , 因为当两者相等时必然存在以下其中一种情况

①、对于某个数 pk , 既不存在 a - pk , 也不存在 b - pk

②、对于某个数 pk , 它既需要存在于集合 A ，也需要存在于集合 B

如果 find（rootA）!= find（rootB），则对应输出就可以了

AC_Coder

#include<bits/stdc++.h>
#define rep(i,a,n) for (int i=a;i<=n;i++)
#define int long long
using namespace std;
const int N = 2e5 + 10;
int far[N] , p[N] , n , a , b;
map<int , int>vis;
int find(int x)
{
    if(x == far[x]) return x;
    return far[x] = find(far[x]);
}
void Union(int x , int y)
{
    int tx = find(x) , ty = find(y);
    if(tx != ty) far[tx] = ty;
}
signed main()
{
    cin >> n >> a >> b;
    rep(i , 1 , n) cin >> p[i] , vis[p[i]] = i , far[i] = i;
    int rootA = n + 1 , rootB = n + 2;
    far[rootA] = n + 1 , far[rootB] = n + 2;
    rep(i , 1 , n)
    {
        if(vis[a - p[i]]) Union(i , vis[a - p[i]]);
        else Union(i , rootB);
        if(vis[b - p[i]]) Union(i , vis[b - p[i]]);
        else Union(i , rootA);    
    } 
    if(find(rootA) == find(rootB)) return cout << "NO\n" , 0;
    cout << "YES\n";
    rep(i , 1 , n) 
    {
        if(find(rootA) == find(i)) 
            cout << 0 << " "; 
        else 
            cout << 1 << " ";
    }
    cout << '\n'; 
    return 0;
}