kmjp's blog

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

Codeforces #423 Div1 E. Rusty String

本番これよく解けたな。
http://codeforces.com/contest/827/problem/E

問題

1文字のワイルドカード'?'および'V''K'の3種類の文字で構成されるN文字の文字列Sがある。
ワイルドカードの中身を好きに決められる場合、Sを右にk文字シフトしても元のSと重なる(言い換えればSのprefix k文字とsuffix k文字が一致する)ようなkを列挙せよ。

解法

?の扱いが面倒なので、まずSと重ならない条件を考える。
S[i]とS[i+m]が重ならない、すなわちどちらかがVでどちらかがWだったとする。
この場合、kがmの約数の場合、S[i],S[i+k],S[i+2k],...,S[i+m]が一致しなければならないので、間の文字種別はともかくkも重ならない。

よって、?は無視してS[i]とS[i+m]が重ならないようなmさえ列挙できれば、どのmの約数にもならないkが解の集合となる。
あとはS[i]とS[i+m]が重なるかどうかである。

このようなマッチング処理はFFTで解くことができる。
以下の2つの多項式を考える。

  • P(x) = sum_i(a_i * x^i) (a_iはS[i]がVなら1、それ以外は0)
  • Q(x) = sum_i(b_i * x^i) (b_iはS[N-1-i]がKなら1、それ以外は0)

P(x)*Q(x)=R(x) = sum_i(c_i * x^i)を考える。
c_iはm=N-1-jとしたときにシフト前のS[i]=Vとシフト後のS[i+m]=Kで重なるiの個数を示す。
これが0であるかどうかで、m文字シフト時に重なるかどうかが判定できる。

VとKを逆にしたパターンも同様に判定するとよい。

typedef complex<double> Comp;

vector<Comp> fft(vector<Comp> v, bool rev=false) {
	int n=v.size(),i,j,m;
	
	for(i=0,j=1;j<n-1;j++) {
		for(int k=n>>1;k>(i^=k);k>>=1);
		if(i>j) swap(v[i],v[j]);
	}
	for(int m=2; m<=n; m*=2) {
		double deg=(rev?-1:1) * 2*acos(-1)/m;
		Comp wr(cos(deg),sin(deg));
		for(i=0;i<n;i+=m) {
			Comp w(1,0);
			for(int j1=i,j2=i+m/2;j2<i+m;j1++,j2++) {
				Comp t1=v[j1],t2=w*v[j2];
				v[j1]=t1+t2, v[j2]=t1-t2;
				w*=wr;
			}
		}
	}
	if(rev) FOR(i,n) v[i]*=1.0/n;
	return v;
}

vector<Comp> MultPoly(vector<Comp> P,vector<Comp> Q) {
	P=fft(P), Q=fft(Q);
	for(int i=0;i<P.size();i++) P[i]*=Q[i];
	return fft(P,true);
}

int T,N;
string S;
int NG[505050];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>T;
	while(T--) {
		cin>>N>>S;
		
		int L=1;
		while(L<=2*N) L*=2;
		vector<Comp> VA(L,0),KA(L,0),VB(L,0),KB(L,0);
		FOR(i,N) {
			if(S[i]=='V') VA[i]=VB[N-1-i]=1;
			if(S[i]=='K') KA[i]=KB[N-1-i]=1;
		}
		auto A=MultPoly(VA,KB);
		auto B=MultPoly(KA,VB);
		
		FOR(i,N+1) NG[i]=0;
		int cnt=1;
		for(i=N-1;i>=1;i--) {
			if((int)(A[N-1-i].real()+0.2)>0 || (int)(B[N-1-i].real()+0.2)>0) NG[i]=1;
			for(j=i*2;j<=N;j+=i) NG[i] |= NG[j];
			if(NG[i]==0) cnt++;
		}
		
		_P("%d\n",cnt);
		for(i=1;i<=N;i++) if(NG[i]==0) _P("%d%c",i,(i==N)?'\n':' ');
		
	}
}

まとめ

なんか最近こういう01マッチングにFFT使う問題好き。