【CF468B】Two Sets

题目

题目链接：https://codeforces.com/problemset/problem/468/B
给出 (n) 个各不相同的数字，将它们分别放入 (A) 和 (B) 两个集合中，使它们满足：
若数字 (x) 在集合 (A) 中，那么数字 (a-x) 也在集合 (A) 中；
若数字 (x) 在集合 (B) 中，那么数字 (b-x) 也在集合 (B) 中。

思路

考虑 2-sat。将 (x) 拆成两个点 (x_A) 和 (x_B)。
如果 (x) 和 (y) 和为 (A)，那么将 (x_A) 与 (y_A) 连边，(y_B) 与 (x_B) 连边。
如果不存在数字 (A-x)，那么从 (x_A) 向 (x_A) 连边。
然后跑 tarjan 即可。
时间复杂度 (O(n+m))。

代码

#include <bits/stdc++.h>
using namespace std;

const int N=200010;
int Q,n,A,B,tot,cnt,a[N],head[N],dfn[N],low[N],col[N];
bool flag,vis[N];
unordered_map<int,int> pos; 
stack<int> st;

int read()
{
	int d=0; char ch=getchar();
	while (!isdigit(ch)) ch=getchar();
	while (isdigit(ch)) d=(d<<3)+(d<<1)+ch-48,ch=getchar();
	return d;
}

struct edge
{
	int next,to;
}e[N*2];

void add(int from,int to)
{
	e[++tot].to=to;
	e[tot].next=head[from];
	head[from]=tot;
}

void tarjan(int x)
{
	dfn[x]=low[x]=++tot; vis[x]=1;
	st.push(x);
	for (int i=head[x];~i;i=e[i].next)
	{
		int v=e[i].to;
		if (!dfn[v])
		{
			tarjan(v);
			low[x]=min(low[x],low[v]);
		}
		else if (vis[v])
			low[x]=min(low[x],dfn[v]);
	}
	if (dfn[x]==low[x])
	{
		cnt++;
		int y;
		do {
			y=st.top(); st.pop();
			col[y]=cnt; vis[y]=0;	
		} while (y!=x);
	}
}

int main()
{
	memset(head,-1,sizeof(head));
	n=read(); A=read(); B=read();
	for (int i=1;i<=n;i++)
	{
		a[i]=read();
		pos[a[i]]=i;
	}
	for (int i=1;i<=n;i++)
	{
		if (pos.find(A-a[i])==pos.end()) add(i,i+n);
			else add(i,pos[A-a[i]]),add(pos[A-a[i]]+n,i+n);
		if (pos.find(B-a[i])==pos.end()) add(i+n,i);
			else add(i+n,pos[B-a[i]]+n),add(pos[B-a[i]],i);
	}
	tot=0; flag=1;
	for (int i=1;i<=2*n;i++)
		if (!dfn[i]) tarjan(i);
	for (int i=1;i<=n;i++)
		if (col[i]==col[i+n])
		{
			flag=0;
			break;
		}
	if (flag) printf("YES
");
		else return printf("NO
"),0;
	for (int i=1;i<=n;i++)
		printf("%d ",col[i]>col[i+n]);
	return 0;
}