Codeforces #268 Div1 B. Two Sets

意外に手こずった。
http://codeforces.com/contest/468/problem/B

問題

N要素のdistinctな整数集合P[i]及び整数a,bがある。
この集合を２つの集合A,Bに分割したい。この時以下を満たす分割があれば答えよ。

• 整数xがAに含まれるなら、a-xもAに含まれる。
• 整数xがBに含まれるなら、b-xもBに含まれる。

解法

まずA==Bの時はコーナーケースとして、全要素をAに入れてよいか検証するだけ。
以下はA!=Bを前提とする。

本番に思いついたのは、「P[i]とP[j]は合わせてAまたはBに入れられる」という条件をもとにグラフを作り二部マッチングをすることだった。
これはアルゴリズムとしてはあっているが、Nが割と大きいのでTLEする。

よってより高速な手段を考える必要がある。
少し考えてみると、全要素が「AかBどちらに入るかわからない」ということはなく、一部の要素xはAかBどちらにしか入らない(a-xとb-xのどちらかがP中にない)ことがわかる。
そこで、そのような選択肢が限られているものから順にA,Bどちらに入るか確定させていけばよい。

自分自身と対、すなわちP[i]*2==AとかP[i]*2==Bとなるケースもあるので注意。

ll N,A,B;
ll P[100005];
map<int,int> M;
int R[100005];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>A>>B;
	
	FOR(i,N) cin>>P[i], M[P[i]]=i;
	
	if(A==B) {
		FOR(i,N) if(M.count(A-P[i])==0) return _P("NO\n");
		_P("YES\n");
		FOR(i,N) _P("0 ");
		return _P("\n");
	}
	
	queue<int> F;
	FOR(i,N) {
		R[i] = (M.count(A-P[i])<<0)+(M.count(B-P[i])<<1);
		if(R[i]==1 || R[i]==2) F.push(i);
		if(R[i]==0) return _P("NO\n");
	}
	
	while(F.size()) {
		int cur=F.front();
		F.pop();
		
		if(R[cur]==1) {
			if(M.count(A-P[cur])==0) return _P("NO\n");
			i=M[A-P[cur]];
			if(R[i]==2) return _P("NO\n");
			if(R[i]==3) R[i]=1, F.push(i);
			if(M.count(B-P[cur])) {
				j=M[B-P[cur]];
				if(R[j]==2) return _P("NO\n");
				if(R[j]==3) R[j]=1, F.push(j);
			}
		}
		if(R[cur]==2) {
			if(M.count(B-P[cur])==0) return _P("NO\n");
			i=M[B-P[cur]];
			if(R[i]==1) return _P("NO\n");
			if(R[i]==3) R[i]=2, F.push(i);
			if(M.count(A-P[cur])) {
				j=M[A-P[cur]];
				if(R[j]==1) return _P("NO\n");
				if(R[j]==3) R[j]=2, F.push(j);
			}
		}
	}
	
	FOR(i,N) if(R[i]!=1 && R[i]!=2) {
		if(P[i]*2==A) R[i]=1;
		if(P[i]*2==B) R[i]=2;
		if(R[i]!=1 && R[i]!=2) return _P("NO\n");
	}
	_P("YES\n");
	FOR(i,N) _P("%d%s", R[i]-1,(i<N-1)?" ":"\n");
}

まとめ

本番二部マッチングでいいの思いついたと思ったけど、周り皆もっとシンプルなコードだったのよね。