Codeforces 268B Two Sets

题目链接:http://codeforces.com/contest/468/problem/B

#include<iostream>
#include<map>
using namespace std;

int N, A, B;
int g[100100], P[100100];

bool must_be_in_A[100100];
bool must_be_in_B[100100];

map<int,int> index_of;

// standard DSU
int find(int x)
{
  if (x == g[x]) 
  {
    return x;
  }
  else
  {
    g[x] = find(g[x]);
    return g[x];
  }
}
int connect(int x, int y)
{
  x = find(x);
  y = find(y);
  g[x]=y;
}

// Note: 
// a-inverse of x means A-x 
// b-inverse of x means B-x 

int main()
{
  ios_base::sync_with_stdio(0);
  cin >> N >> A >> B;
  for (int i = 1; i <= N; i++)
  {
    cin >> P[i];
    index_of[P[i]] = i;
  }

  for (int i = 1; i <= N; i++) g[i] = i;

  // check if there's any element 
  // that has neither a-inverse nor b-inverse
  // if there's one, it cannot be in either A or B
  // thus the answer is "NO"

  for (int i = 1; i <= N; i++)
  {
    if (index_of[A-P[i]]==0 && index_of[B-P[i]] == 0)
    {
      cout << "NO" << endl;
      return 0;
    }
  }

  // As shown in editorial
  // x must be in the same set as
  // its a-inverse and b-inverse (if they exist)
  // thus x must be connected to a-inverse and b-inverse 
  for (int i = 1; i <= N; i++)
  {
    int ainv = A-P[i];
    if (index_of[ainv] >= 1)
    {
      connect(i, index_of[ainv]);
    }
    int binv = B-P[i];
    if (index_of[binv] >= 1)
    {
      connect(i, index_of[binv]);
    }
  }

  // if a connected set contains element which cannot be in B
  // (an element without b-inverse)
  // then every element in the connected set must be in A, and vice versa.
  // Here we will keep the information on "root" node in DSU.

  for (int i = 1; i <= N; i++)
  {
    // whether each inverse exists;
    bool ainv = index_of[A-P[i]] != 0;
    bool binv = index_of[B-P[i]] != 0;
    if (ainv && !binv)
    {
      must_be_in_A[find(i)] = true;
    }
    if (binv && !ainv)
    {
      must_be_in_B[find(i)] = true;
    }
  }

  // if there is any set which contains 
  // both an element that must be in A
  // and an element that must be in B
  // then it's impossible
  for (int i = 1; i <= N; i++)
  {
    if (must_be_in_A[i] && must_be_in_B[i])
    {
      cout << "NO" << endl;
      return 0;
    }
  }

  // Finally, we found the answer!
  cout << "YES" << endl;

  for (int i = 1; i <= N; i++)
  {
    int which_set = -1;
    if (must_be_in_A[find(i)])
    {
      which_set = 0;
    }
    else if (must_be_in_B[find(i)])
    {
      which_set = 1;
    }
    else
    {
      // if we can do either way, just say set A.
      which_set = 0;
    }
    cout << which_set << (i == N ? '\n' : ' ');
  }
  return 0;
}
点赞