Codeforces #930 : Div1 B. Pinball

ちょっと出来の悪かった回。
https://codeforces.com/contest/1936/problem/B

問題

左右にNマスが並んでおり、各マス右矢印か左矢印が並んでいる。
どこかのマスにボールを置くと、以下のように状態が変わる。

• ボールがマスの矢印の向きに1つ移動する。
• その後、元のマスの矢印の向きが反転する。

各マスにボールを置いた時それぞれについて、ボールがNマス外に出るまでに何マス移動するか答えよ。

解法

ボールの移動方向が反転する回数は、初期位置の左側にある右向きの矢印と、右側にある左向きの矢印で決まる。
また、移動距離はそれらの座標で決まる。

両矢印のマスの座標の一覧と、座標の累積和をもっておけば、上記はO(logN)で計算できる。

int T,N;
string S;
vector<int> L,R;
vector<ll> LS,RS;

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>T;
	while(T--) {
		cin>>N>>S;
		L.clear();
		R.clear();
		LS={0LL};
		RS={0LL};
		FOR(i,N) {
			if(S[i]=='<') {
				L.push_back(i);
				LS.push_back(LS.back()+i);
			}
			else {
				R.push_back(i);
				RS.push_back(RS.back()+i);
			}
		}
		
		FOR(i,N) {
			int nr=lower_bound(ALL(R),i+1)-R.begin();
			int nl=L.end()-lower_bound(ALL(L),i);
			ll ret=-1;
			if(S[i]=='<') {
				nl=min(nl,nr+1);
				nr=min(nr,nl);
				if(nl==nr) {
					x=lower_bound(ALL(L),i)-L.begin();
					ret=N+i+2*(LS[x+nl]-LS[x+1]);
					x=lower_bound(ALL(R),i)-R.begin();
					ret-=2*(RS[x]-RS[x-nr]);
				}
				else {
					x=lower_bound(ALL(L),i)-L.begin();
					ret=i+2*(LS[x+nl]-LS[x+1])+1;
					x=lower_bound(ALL(R),i)-R.begin();
					ret-=2*(RS[x]-RS[x-nr]);
				}
				
			}
			else {
				nr=min(nr,nl+1);
				nl=min(nl,nr);
				if(nl==nr) {
					x=lower_bound(ALL(L),i)-L.begin();
					ret=2*(LS[x+nl]-LS[x]);
					x=lower_bound(ALL(R),i+1)-R.begin();
					ret-=2*(RS[x]-RS[x-nr])-R[x-1];
					ret++;
				}
				else {
					x=lower_bound(ALL(L),i)-L.begin();
					ret=N+2*(LS[x+nl]-LS[x]);
					x=lower_bound(ALL(R),i+1)-R.begin();
					ret-=2*(RS[x]-RS[x-nr])-R[x-1];
				}
			}
			cout<<ret<<" ";
		}
		cout<<endl;
		
	}
}

まとめ

これそんなに難しく見えないけど本番なんか苦戦してるな…。