【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\n");
		else return printf("NO\n"),0;
	for (int i=1;i<=n;i++)
		printf("%d ",col[i]>col[i+n]);
	return 0;
}