kmjp's blog

競技プログラミング参加記です

Codeforces #162 Div1 D. Colorful Stones

SRMのDiv1 Med埋めも一段落ということで、今度は初参加以降のCodeforces Div1 D埋めをしていきます。
http://codeforces.com/contest/264/problem/D

問題

1列に並んだ石が2組ある。
それぞれの石はR,G,Bのいずれかの色である。
初期状態ではそれぞれの先頭の石に1個ずつコマがある。

ここで以下の処理を繰り返す。
プレイヤーがR,G,Bのいずれかの色を選択すると、両列のコマがそれぞれその色の石に乗っている場合、そのコマを1つ前に進める。

石の色の並びが文字列S,Tで与えられている。
初期状態から遷移できるコマの位置の対の数を答えよ。

解法

基本的には尺取法で行く。
すなわち、1列目で位置xにいるとき、2列目のいる位置の下限a・上限bを求めて(b-a+1)の和を取ればよい。

まずxに対するa,bの求め方を考える。

  • 下限aは、最短の処理で1列目でxに到達するように処理した場合、すなわちS[0]~S[x-1]の色を順に選択したときの2列目の位置となる。
  • 上限は、逆に2列目で最短処理でT[b+1]に到達しようとして、1列目でx+1番目に到達してしまうようなbを求めればよい。

いずれも尺取法の要領で全部でO(|S|+|T|)で処理できる。

ただし、1列目で位置xでありながら、2列目の位置a~bの間で到達できない位置yがある。
それはT[y]直前の2文字T[y-2]・T[y-1]の並びがS[x]直前の2文字S[x-2]・S[x-1]の逆となっている場合である。

よって事前にT[y]の直前2文字T[y-2]・T[y-1]の登場パターンを累積和で計算しておき、上記尺取法の過程で(b-a+1)から引けばよい。

string S,T;
int LS,LT;
int AP[1000002],BP[1000002];
int mat[1000002][3][3];

void solve() {
	int f,i,j,k,l,x,y;
	cin>>S>>T;
	LS=S.size();
	LT=T.size();
	FOR(i,LS) S[i]=(S[i]=='R')?0:(S[i]=='G')?1:2;
	FOR(i,LT) T[i]=(T[i]=='R')?0:(T[i]=='G')?1:2;
	
	for(i=1;i<LT;i++) {
		memmove(mat[i+1],mat[i],sizeof(mat[i]));
		mat[i+1][T[i]][T[i-1]]++;
	}
	
	x=y=0;
	FOR(i,LS) AP[i]=x, x+=(x<LT&&S[i]==T[x]);
	FOR(i,LT) BP[i]=y, y+=(y<LS&&S[y]==T[i]);
	
	ll ret=0;
	y=0;
	FOR(i,LS) {
		x=AP[i];
		while(y<LT && i>=BP[y]) y++;
		
		ret+=y-x;
		if(i>0 && S[i-1]!=S[i]) ret -= mat[y][S[i-1]][S[i]]-mat[x][S[i-1]][S[i]];
	}
	cout << ret << endl;
}

まとめ

尺取法で行くことは思いついたけど、到達できないケースに最後2文字が逆となるパターンがある、ということに思い至らなかった…。