SRM 286 DIV1 Hard InfiniteSoup

問題 Editorial

問題

無限のグリッドにそれぞれ小文字アルファベットが書かれている。
(i,j)に書かれている文字は(g[i mod R][j mod C] where gはR*Cのサイズ)である。
(0,0)から(x,y) (x,yは非負整数)を通る半直線上の(厳密にその座標を通った)文字を繋げると無限の文字列ができる。
整数kが与えられる。[(x,y) | x <- [0..k], y <- [0..k]]で表される無限の文字列のうち、wordsのそれぞれに一致するものの数をvectorで返せ。
ただし、同じ半直線(例えば(1,1)と(2,2))は重複して数えない

  • 1 <= R, C <= 35
  • 1 <= |words| <= 25
  • 1 <= |words[i]| <= 50
  • 1 <= k <= 200

解答

まず、半直線(x,y)は(x mod R, y mod C)で同じなので、(x < R, y < C)の場合を考えてあとでその数だけ掛ければいい。
(x, y)は(gcd(x, y) = 1)の時のみ考えれば重複しない。
このことを考えると、文字列の数は最大でもたったの721個以下。十分列挙・それぞれで判定ができる。
さて、無限の文字列中の検索はどのようにすればいいか?これは単純で、適当な長さ(50文字以上)になるまで文字列を繰り返したあと文字列を2倍すれば、有限の文字列で判定ができる。
文字列の長さの最長は(35*35*2 = 2450)である。各単語を全てO(|文字列| + |単語|)程度で検索すれば(721*25*2450 ≒ 4*10^7)程度で、十分間に合う。これにはKMPなどが使える。
また、自分はAho-Corasickを用いて全ての単語を一気に検索した

コード

自分のこのコードでは(x < R, y < C)に気づかなくて24465の文字列を全てやってる。でもAho-Corasickを用いたら間に合った(0.7s程度)。
Aho-Corasickは最初 Spaghetti Source のものを多少改変してそのまま貼ったらバグった。そこに書いてあったスライドを参考に一行足したら通った。これは Spaghetti Source のコードが間違っているということなのか?ともかく今バグを踏んでおいてよかった

#include <string>
#include <vector>
#include <cstring>
#include <queue>
#define rep(i,n) for(int (i)=0;(i)<(int)(n);++(i))
#define rer(i,l,u) for(int (i)=(int)(l);(i)<=(int)(u);++(i))
#define mset(m,v) memset(m,v,sizeof(m))
using namespace std;
typedef vector<int> vi;  
typedef long long ll;   
typedef vector<string> vs; 

const int Alphabets = 26;
struct PMA {
	PMA *fail;
	PMA *next[Alphabets];
	ll accept;
	PMA(): fail(NULL), accept(0) { mset(next, 0); }
};
PMA *buildPMA(const vs &p) {
	PMA *root = new PMA;
	for (int i = 0; i < p.size(); ++i) { // make trie
		PMA *t = root;
		for (int j = 0; j < p[i].size(); ++j) {
			char c = p[i][j] - 'a';
			if (t->next[c] == NULL) t->next[c] = new PMA;
			t = t->next[c];
		}
		t->accept |= 1LL << i;
	}
	queue<PMA*> Q;
	rep(c, Alphabets) {
		if (root->next[c]) {
			root->next[c]->fail = root;
			Q.push(root->next[c]);
		} else root->next[c] = root;	//ここはスライドのにはなかったけど…
	}
	while (!Q.empty()) {
		PMA *t = Q.front(); Q.pop();
		rep(c, Alphabets) {
			if (t->next[c]) {
				Q.push(t->next[c]);
				PMA *r = t->fail;
				while (!r->next[c]) r = r->fail;
				t->next[c]->fail = r->next[c];
				t->next[c]->accept |= r->next[c]->accept;	//spaghetti source のものではここがなかったけど…
			}
		}
	}
	return root;
}
ll match(char *t, int n, PMA *v) {
	ll r = 0;
	for (int i = 0; i < n; ++i) {
		char c = t[i] - 'a';
		while (!v->next[c]) v = v->fail;
		v = v->next[c];
		r |= v->accept;
	}
	return r;
}

template<typename T>T gcd(T x, T y) { return y == 0 ? x : gcd(y,x%y); }
struct InfiniteSoup {
	vector <int> countRays(vector <string> g, vector <string> words, int k) {
		int R = g.size(), C = g[0].size();
		int n = words.size();
		PMA *pma = buildPMA(words);
		vi r(n, 0);
		rer(y, 0, k) rer(x, 0, k) if(gcd(y, x) == 1) {
			int i = 0, j = 0;
			static char s[10000];
			int l = 0;
			do{
				s[l ++] = g[i][j];
				(i += y) %= R, (j += x) %= C;
			}while(i!=0||j!=0||l<=100);
			memcpy(s+l, s, l); l *= 2;
			s[l] = 0;
			ll t = match(s, l, pma);
			rep(w, n) if(t >> w & 1) r[w] ++;
		}
		return r;
	}
};