CSAcademy Round #12 : E. Prefix Suffix Counting

桁DP2連発な感じでぐったり。
https://csacademy.com/contest/round-12/#task/prefix-suffix-counting

問題

２つの巨大な数N,Mがある。
1≦x≦Nである|M|桁以上のxのうち、上位M桁と下位M桁がMと一致するものの個数を答えよ。

解法

x	がNの場合とN未満の場合でそれぞれ考える。

• |x|がM以上2M未満の場合
  • 上位M桁と下位M桁がわかるとxが1つ確定する。
  • 上位M桁と下位M桁が一部重複する場合、重複部分が矛盾を起こさないか判定する。自分はZ-algorithmで求めた。ローリングハッシュで間に合うかは未確認。
• |x|が2M以上の場合
  • |x|が|N|未満の場合
    • |x|-(2M)桁は0～9のどれが入ってもいいので、10^(|x|-2M)を解に計上する。
  • |x|が|N|の場合
    • 上位・下位M桁が確定した整数のうち、N以下のものを桁DPの要領で求めよう。

string N,M;
int NL,ML;
ll p10[1010101];
ll mo=1000000007;

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);
		}
	}
	return v;
}

ll dfs(int id,string& N,string& T,int nx) {
	if(id==N.size()) return 1;
	
	if(T[id]=='X') {
		return ((N[id]-'0')*p10[nx-1] + dfs(id+1,N,T,nx-1))%mo;
	}
	else {
		if(N[id]>T[id]) return p10[nx];
		if(N[id]<T[id]) return 0;
		return dfs(id+1,N,T,nx);
	}
	
	
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>M;
	NL=N.size();
	ML=M.size();
	
	p10[0]=1;
	FOR(i,1010000) p10[i+1]=p10[i]*10%mo;
	vector<int> Z=Zalgo(M);
	
	ll ret=0;
	for(int i=ML;i<NL;i++) {
		if(i<2*ML) {
			j=i-ML;
			if(Z[j]>=ML-j) ret++;
		}
		else {
			ret += p10[i-(2*ML)];
		}
	}
	
	
	if(NL>=ML) {
		string T(NL,'X');
		FOR(i,ML) T[i]=M[i];
		FOR(i,ML) {
			j=NL-ML+i;
			if(T[j]=='X') T[j]=M[i];
			if(T[j]!=M[i]) break;
		}
		if(i==ML) ret += dfs(0,N,T,count(T.begin(),T.end(),'X'));
	}
	cout<<ret%mo<<endl;
	
}

まとめ

桁DP面倒で好きじゃない…。