kmjp's blog

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

TopCoder SRM 797 : Div1 Medium RollMe

ライブラリ不足。
https://community.topcoder.com/stat?c=problem_statement&pm=16757&rd=18435

問題

N面ダイスがあり、各面には0~(N-1)の数値が書かれている。
また、各面cが出る確率P(c)が与えられる。

ここで文字列Gが与えられる。
ダイスを繰り返し振り、出た目を並べた文字列を考える。
末尾の|G|文字が初めてGに一致するまでにかかるダイスを振る回数の期待値を求めよ。

解法

f(i) := 末尾がi文字一致した状態で、(i+1)文字一致した状態にかかるまでに必要なダイスの回数の期待値
g(i) := f(i)の累積和
とする。求めたいのはf(0)+f(1)+....+f(|G|-1)であり、これはg(|G|-1)に等しい。

では、f(i)を求めることを考える。
最初にAho-CorasickなりKMP法なりで、N面の各文字cが出たときの遷移先T(c)を求めておこう。
N面のうち(i+1)字一致した状態に遷移できるのは1通りだけで、それ以外は遷移を進めることができない。
その場合、またi文字一致した状態に戻るためf(T(c))+f(T(c)+1)+f(T(c)+2)+....+f(i-1) = g(i-1)-g(T(c)-1)だけかかる。
そこでT(c)=i+1となるcをc'とすると
f(i) = 1 + sum(P(c)*(f(i)+g(i-1)-g(T(c)-1)) (ただしc'は除く)
となるので、
f(i) = (1 + sum(P(c)*(g(i-1)-g(T(c)-1))) / (1+sum(P(c)) (ただしc'は除く)
と変形できる。

const ll mo=1000000007;
ll dp[5050],S[5050];

struct KMP {
	vector<int> failto;
	string p;
	KMP(string s) :p(s) {
		int n=s.size();
		failto.resize(n+1);
		int i,j; j=failto[0]=-1;
		for(i=1;i<=n;i++) {
			while(j>=0&&s[j]!=s[i-1]) j=failto[j];
			failto[i]=++j;
		}
	}
	pair<int,int> match(string& s) { //sのprefixのサイズと、マッチした長さ
		int a,b;
		for(a=b=0;a<s.size();a++) {
			while(b>=0&&p[b]!=s[a]) b=failto[b];
			b++;
			if(b==p.size()) return {a+1,b};
		}
		return {s.size(),b};
		
	}
};


ll modpow(ll a, ll n = mo-2) {
	ll r=1;a%=mo;
	while(n) r=r*((n%2)?a:1)%mo,a=a*a%mo,n>>=1;
	return r;
}

class RollMe {
	public:

	int solve(vector <int> die, string goal) {
		vector<ll> D;
		int N=die.size(),M=goal.size();
		int i,j;
		int s=0;
		
		FORR(d,die) s+=d;
		FORR(d,die) D.push_back(d*modpow(s)%mo);
		
		FORR(g,goal) g-='0';
		KMP kmp(goal);
		
		FOR(i,M) {
			dp[i]=1;
			ll same=0;
			FOR(j,N) {
				string s=goal.substr(0,i)+((char)j);
				int tar=kmp.match(s).second;
				if(tar<=i) {
					same+=D[j];
					(dp[i]+=D[j]*(mo+S[i]-S[tar]))%=mo;
				}
			}
			(dp[i]=dp[i]*modpow(1+mo-same%mo))%=mo;
			S[i+1]=(S[i]+dp[i]+mo)%mo;
		}
		
		return S[M];
	}
}

まとめ

本番、KMP法のライブラリが手元にないので探して実装したら、参照先が誤っていたというドタバタが発生。