kmjp's blog

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

Codeforces #248 Div1 C. Tachibana Kanade's Tofu

Codeforces

まーた面倒な問題。
http://codeforces.com/contest/434/problem/C

問題

M進数で２つの値L,Rが与えられる。
また、同じくM進数の文字列A[i]とポイントV[i]が計N個与えられる。

ある値xのスコアは、M進数表記においてA[i]を含むたびにポイントV[i]を加算し、その合計を取ったものである。
L～Rの値のうち、(leading zeroを除き)スコアがK以下のものを答えよ。

解法

Editorialを見て解答。
探索する文字列が複数あるので、Aho-Corasick法を用いる。
A[i]からAho-Corasickの状態遷移図を作成し、各状態において加算されるポイントを求める。

L～Rのうち条件を満たす数を求めるので、1～R中の数から1～(L-1)中の数を引けばよい。
1～Xのうち条件を満たす数を求めるには、状態として[処理済みの桁数][状態遷移図中の位置][取得済み総ポイント][その桁で任意の値を利用可能か][最上位桁か]を持って、遷移図中で桁DPしていけばよい。

Aho-Corasickの状態数はA[i]の総長ALであり、全計算量はO((logR)*AL*K*M)である。

int N,M,K;
ll mo=1000000007;
int L,R;
string LL,RR;
vector<string> S;
int V[201];

vector<vector<int> > TR;
vector<int> AC;

ll dpdp[2][201][501][2][2];

int find(string s) {
	int cur=0;
	ITR(it,s) if((cur=TR[cur][*it])==0) return -1;
	return cur;
}
void create(vector<string> S) {
	int i;
	TR.push_back(vector<int>(256));
	// trie
	ITR(it,S) {
		int cur=0;
		ITR(c,(*it)) {
			if(TR[cur][*c]==0) TR.push_back(vector<int>(256)),TR[cur][*c]=TR.size()-1;
			cur=TR[cur][*c];
		}
		AC.resize(TR.size());
		AC[cur]+=V[it-S.begin()];
	}
	// failure
	queue<int> Q;
	FOR(i,255) if(TR[0][i]) TR[TR[0][i]][255]=0, Q.push(TR[0][i]);
	
	while(!Q.empty()) {
		int k=Q.front(); Q.pop();
		FOR(i,255) if(TR[k][i]) {
			Q.push(TR[k][i]);
			int pre=TR[k][255];
			while(pre!=0 && TR[pre][i]==0) pre=TR[pre][255];
			TR[TR[k][i]][255]=TR[pre][i];
			AC[TR[k][i]] += AC[TR[pre][i]];
		}
	}
}

ll dodo(string SS) {
	ll ret=0,rr;
	int i,j,k,x,y,pos,val,b1,b2;
	
	ZERO(dpdp);
	dpdp[0][0][0][1][1]=1;
	FOR(i,SS.size()) {
		int cur=i%2,tar=cur^1;
		ZERO(dpdp[tar]);
		FOR(b1,2) FOR(b2,2) FOR(val,K+1) FOR(pos,TR.size()) if((rr = dpdp[cur][pos][val][b1][b2])) {
			FOR(j,M) {
				if(b1 && j>SS[i]) break;
				y=(b1 && j==SS[i]);
				if(b2&&(j==0)) {
					dpdp[tar][pos][val][y][1] += rr;
					dpdp[tar][pos][val][y][1] %= mo;
					continue;
				}
				x=pos;
				while(x && TR[x][j]==0) x=TR[x][255];
				x=TR[x][j];
				if(val+AC[x]<=K) {
					dpdp[tar][x][val+AC[x]][y][0] += rr;
					dpdp[tar][x][val+AC[x]][y][0] %= mo;
				}
			}
		}
	}
	x=SS.size()%2;
	FOR(i,201) FOR(j,K+1) ret+=dpdp[x][i][j][0][0]+dpdp[x][i][j][0][1]+dpdp[x][i][j][1][0]+dpdp[x][i][j][1][1];
	return ret%mo;
}

void solve() {
	int f,i,j,k,l,x,y;
	
	cin>>N>>M>>K;
	cin>>L;
	FOR(i,L) cin>>x, LL+=x;
	cin>>R;
	FOR(i,R) cin>>x, RR+=x;
	
	FOR(i,N) {
		cin>>j;
		string s;
		while(j--) cin>>x, s+=x;
		S.push_back(s);
		cin>>V[i];
	}
	
	create(S);
	
	i=LL.size()-1;
	LL[i]-=1;
	while(LL[i]<0) LL[i]+=M, LL[--i]-=1;
	while(LL.size()>0 && LL[0]==0) LL.erase(LL.begin());
	
	cout << (dodo(RR)-dodo(LL)+mo)%mo << endl;
}

まとめ

Aho-Corasick法を知らなかったので頑張ってライブラリを作ったら、単なるマッチング判定じゃなく全状態遷移におけるDPが必要でぐったり。