TopCoder SRM 853 : Div1 Medium SubstringCoverage

しょうもないミスで落とした…。
https://community.topcoder.com/stat?c=problem_statement&pm=18168

問題

文字列TがSでカバーされるとは、Tの各文字が、T中に部分文字列として現れるSの一部となることを意味する。

文字列Needleと整数Hが与えられる。
H文字以下の文字列のうち、Needleでカバーされる物について、そのような文字列長の個数はいくつか。

解法

初期状態でX=Needleとする。
XにNeedleのsuffixの何文字かを追加したとき、その文字列もNeedleでカバーされるような文字数の一覧は、Z-algorithmなどでsん出できる。

そのような文字数の集合をAとすると、解は|Needle|にAの要素を何個でも繰り返し加算して作れるH以下の整数の個数となる。
これを数え上げよう。

f(x) := |X| % |Needle| = xとなるXのうち、Needleでカバーされる最短のXの長さ
とする。
f(0)～f(|Needle-1|)をダイクストラ法の要領で求めることができる。

int D[5050];

class SubstringCoverage {
	public:
	int count(string needle, int H) {
		int N=needle.size();
		vector<int> V={N};
		int i;
		for(i=1;i<N;i++) if(needle.substr(0,i)==needle.substr(N-i)) V.push_back(N-i);
		FOR(i,N) D[i]=1<<30;
		D[0]=N;
		priority_queue<pair<int,int>> Q;
		Q.push({-N,0});
		
		while(Q.size()) {
			int co=-Q.top().first;
			int cur=Q.top().second;
			Q.pop();
			if(D[cur]!=co) continue;
			FORR(a,V) if(chmin(D[(cur+a)%N],co+a)) Q.push({-D[(cur+a)%N],(cur+a)%N});
		}
		
		int ret=0;
		FOR(i,N) if(D[i]<=H) ret+=1+(H-D[i])/N;
		
		
		
		return ret;
		
		
	}
}

まとめ

本番中愚直にO(|Needle|^2)要素の配列を作ろうとしてしまった。O(|Needle|)で済ますテク、何度か使ったことあるのに今回思いつけなかった…。