kmjp's blog

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

Codeforces ECR #051: E. Vasya and Big Integers

グダグダかと思ったけどなんか1位だった。
http://codeforces.com/contest/1051/problem/E

問題

非負整数A,L,Rが与えられる。
Aを文字列とみなし、いくつかの部分文字列の連結として表す。
その際、各部分文字列がL以上R以下であるような部分文字列の分割は何通りか。

解法

f(i) := Aの先頭i文字を条件を満たすように分ける分け方
とする。f(0)=1から始めて順次f(i)を求めていき、最終的にf(N)を求めよう。

f(i)がわかっているとき、次に何文字の部分文字列が取れるかを考える。

  • A[i...(i+|L|-1)]≧Lなら|L|文字以上、そうでないなら|L|+1文字以上
  • A[i...(i+|R|-1)]≦Rなら|R|文字以下、そうでないなら|R|-1文字以下

を取ることができる。ここから、f(i+|L|)またはf(i+|L|+1)から、f(i+|R|)またはf(i+|R|-1)までf(i)を足しこむ処理を行う。
これはBITでもいいが、単なる累積和で行う方が計算量が少なくて済む。

問題はA[i...(i+|L|-1)]≧LとA[i...(i+|R|-1)]≦Rの判定である。
自分は最初O(NlogN)のSuffixArrayを使ったところ、A,L,Rが最大10^6桁あることもありTLEした。
SA-ISを持っていないので頭を抱えたが、結局L+AやR+Aに対しZ algorithmを適用すれば大小判定が全体でO(|L|+|R|+|A|)で済む。

vector<int> Zalgo(string s) {
	vector<int> v(1,s.size());
	for(int i=1,l=-1,r=-1;i<s.size();i++) {
		if(i<=r && v[i-l]<r-i+1) v.push_back(v[i-l]);
		else {
			l=i; r=(i>r)?i:(r+1);
			while(r<s.size() && s[r-i]==s[r]) r++;
			v.push_back((r--)-l);
		}
	}
	v.push_back(0);
	return v;
}

string A,L,R,LA,RA;
int NA,NL,NR,LS,RS;;
vector<int> LZ,RZ;

ll mo=998244353;
ll bt[2020202];


void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>A>>L>>R;
	NA=A.size();
	NL=L.size();
	NR=R.size();
	
	LA=L+"@"+A;
	RA=R+"@"+A;
	LZ=Zalgo(LA);
	RZ=Zalgo(RA);
	
	
	bt[0]=1;
	bt[1]=mo-1;
	FOR(i,NA) {
		if(i) (bt[i]+=bt[i-1])%=mo;
		ll now=bt[i];
		if(A[i]=='0') {
			if(L[0]=='0') {
				(bt[i+1]+=now)%=mo;
				(bt[i+2]+=mo-now)%=mo;
			}
			continue;
		}
		
		x=min(NA-i,NL);
		if(x<NL) continue;
		if(LZ[NL+1+i]<x && A[i+LZ[NL+1+i]]<L[LZ[NL+1+i]]) x++;
		y=min(NA-i,NR);
		if(y==NR && RZ[NR+1+i]<y && A[i+RZ[NR+1+i]]>R[RZ[NR+1+i]]) y--;
		
		if(x<=y) {
			(bt[i+x]+=now)%=mo;
			(bt[i+y+1]+=mo-now)%=mo;
		}
		
	}
	cout<<((bt[NA]+bt[NA-1])%mo)<<endl;
}

まとめ

SA-IS、いい加減実装しておこうかな…。