AtCoder AGC #006 : D - Median Pyramid Hard

これは思いつかなかった。
http://agc006.contest.atcoder.jp/tasks/agc006_d

問題

正方形のマスを積み上げたN段のピラミッドを考える。
上からi段目には(2i-1)個のマスからなる。
各段は隙間なく正方形を並べており、各段の中央のマスは左右同じ位置に配置されている。

最下段の各マスに、1～(2N-1)の順列を値を成すよう１つずつ整数を埋めるとする。
最下段より上のマスは、真下かその左右隣りの計３マスの整数のうち中央値を埋めていく。

最下段の各マスに埋められた整数が与えられたとき、最上段のマスに埋められる数値を求めよ。

解法

二分探索で解く。
解がx以上かどうかを判定するためには、最下段の各マスをx未満なら0、x以上なら1、と置き換えたとき最上段が1になるかどうかを判定すればよい。

こうして01に置き換えた場合、上の段がどのように埋まるかを考えてみる。
最下段中央が01交互に並んでいるとする。
その左右それぞれ探索した場合にどちらか01が２マス連続する箇所があった場合、中央に近い側の値が最終的に中央にくる。

よって二分探索の各ステップでは最下段の２連続マスを探すだけなのでO(N)で済む。

int N;
int A[202020];
int B[202020];

int ok(int cur) {
	if(cur>2*N-1) return 0;
	
	int i;
	FOR(i,2*N-1) B[i]=A[i]>=cur;
	int L=N-1,R=N-1;
	while(L>0 && B[L]!=B[L-1]) L--;
	while(R<2*N-1 && B[R]!=B[R+1]) R++;
	return ((N-1)-L<=R-(N-1))?B[L]:B[R];
}


void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,2*N-1) cin>>A[i];
	int cur=1;
	for(i=20;i>=0;i--) if(ok(cur+(1<<i))) cur+=1<<i;
	cout<<cur<<endl;
}

まとめ

終わってみるとコードが異常に短いな…。