square869120Contest #1 : H - 3人の昼食

10億円の昼食を買える中学生が、昼食の値段で喧嘩か…。
http://s8pc-1.contest.atcoder.jp/tasks/s8pc_1_h

問題

N個の正整数A[i]が与えられる。
これを4つの集合W,X,Y,Zと分けたとする。
整数D,Eに対し

• |W|≦E
• |sum(X)-sum(Y)|≦D
• |sum(X)-sum(Z)|≦D

となる分け方は何通りあるか。

解法

公式解説が出る前だったので下記を参考に解きました。
square869120Contest H - 3人の昼食 (The Lunch) - mayoko’s diary

総当たりだとO(N*4^N)かかるので部分点どまり。
そこで半分全列挙を行う。

Nのうち半数に対し総当たりする。
まずはその範囲で|W|及び(sum(X)-sum(Y),sum(X)-sum(Z))の値を数え上げよう。
残り半分も同様に総当たりする。

P(w,a,b) := Nの前半半分について、|W|=w, a=sum(X)-sum(Y), b=sum(X)-sum(Z)となる分け方の組み合わせ数。
Q(w,a,b) := Nの後半半分について以下同様。

あとはw+w'≦E、|a+a'|≦D、|b+b'|≦DとなるようなP(w,a,b)*Q(w',a',b')の総和を取りたい。
Eは小さいのでw,w'は総当たりしてよい。
あとは(a,b)に対し条件を満たす(a',b')の列挙である。
(a,b)に対し、a-D≦a'≦a+D、b-D≦b'≦b+Dを満たす(a',b')におけるP(w,a,b)*Q(w',a',b')を求めたい。

これは平面走査と座標圧縮で対応できる。
まずは以下の整数対を２次元上の座標とみなし、x座標順にソートする。

• P(w,a,b)＞0となる(a,b)
• Q(w',a',b')＞0となる(a',b')に対する(a'-D,b'),(a'+D,b')

あとはソートした頂点に対し順次以下の処理を行う。

• (a'-D,b')を処理する場合、(b'-D)～(b'+D)の区間にQ(w',a',b')を加算する
• (a'+D,b')を処理する場合、(b'-D)～(b'+D)の区間にQ(w',a',b')を減算する
• (a,b)を処理する場合、区間中bにおける整数の総和にP(w,a,b)を掛け答えに合算する

区間に対する加減算が登場するので、実際は登場するb,b'+D,b'-Dを座標圧縮してBITを使うとよい。

int N,M,E;
ll D,A[50];
vector<pair<ll,ll>> B[2][3];

template<class V, int ME> class BIT {
public:
	V bit[1<<ME],val[1<<ME];
	V total(int e) {V s=0;e++;while(e) s+=bit[e-1],e-=e&-e; return s;}
	V add(int e,V v) { val[e++]+=v; while(e<=1<<ME) bit[e-1]+=v,e+=e&-e;}
};
BIT<ll,20> bt;

ll hoge(vector<pair<ll,ll>> S,vector<pair<ll,ll>> T) {
	vector<ll> Y;
	vector<int> L,R;
	
	Y.push_back(-1LL<<60);
	Y.push_back(1LL<<60);
	FORR(r,S) Y.push_back(-D-r.second),Y.push_back(D-r.second);
	FORR(r,T) Y.push_back(r.second);
	sort(ALL(Y));
	Y.erase(unique(ALL(Y)),Y.end());
	int i;
	
	ZERO(bt.bit);
	ZERO(bt.val);
	
	vector<pair<ll,pair<int,ll>>> ev;
	FORR(r,S) ev.push_back({r.first,{1,r.second}});
	FORR(r,T) {
		ev.push_back({-D-r.first,{0,r.second}});
		ev.push_back({D-r.first,{2,r.second}});
	}
	
	ll ret=0;
	sort(ALL(ev));
	FORR(r,ev) {
		if(r.second.first==1) {
			int L=lower_bound(ALL(Y),-D-r.second.second)-Y.begin();
			int R=lower_bound(ALL(Y),D-r.second.second)-Y.begin();
			ret+=bt.total(R)-bt.total(L-1);
		}
		else {
			int y=lower_bound(ALL(Y),r.second.second)-Y.begin();
			if(r.second.first==0) bt.add(y,1);
			else bt.add(y,-1);
		}
	}
	
	return ret;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	int mask;
	cin>>N>>D>>E;
	FOR(i,N) cin>>A[i];
	FOR(mask,1<<(2*min(10,N))) {
		int e=0;
		ll s[3]={};
		FOR(i,min(10,N)) {
			x = (mask>>(2*i))&3;
			if(x==3) e++;
			else s[x]+=A[i];
		}
		if(e<=E) B[0][e].push_back({s[1]-s[0],s[2]-s[0]});
	}
	FOR(mask,1<<(2*max(N-10,0))) {
		int e=0;
		ll s[3]={};
		FOR(i,max(N-10,0)) {
			x = (mask>>(2*i))&3;
			if(x==3) e++;
			else s[x]+=A[i+10];
		}
		if(e<=E) B[1][e].push_back({s[1]-s[0],s[2]-s[0]});
	}
	
	ll ret=0;
	FOR(x,3) FOR(y,3) if(x+y<=E) ret+=hoge(B[0][x],B[1][y]);
	cout<<ret<<endl;
}

まとめ

公式解説の「THE END」が格好いい。