初めて Trie 木を使ったので、ハマったところとかを記事でまとめておきます (完全に自分用)
問題概要
原文参照 → C: たんごたくさん - 天下一プログラマーコンテスト2016本戦(オープンコンテスト) | AtCoder
解説
与えられる単語すべてを検索対象として、Trie 木を作ります。
文字列 の 文字目 を左端とする の部分文字列であって、単語と完全一致するものを考えます。単語の長さが高々 であることから、部分文字列の選び方は 通りしかありません。この制約だと、部分文字列の左端を決め打ちしてこの 通りを全て試すことで間に合うため、それを基にして DP を実装すればよいです。
ただ、部分文字列を作るために substr を使ったり、文字列検索のためにいちいち trie 木の根から走査していくと TLE になるため、前の結果をうまく使う必要があります。自分の実装では、文字列検索時に Trie のポインタを動かして、部分文字列の右端が右に移動する (= 部分文字列が 1 文字増える) ときに前の結果を用いて高速に処理するようにしています。
ソースコード
struct Trie { Trie* node[26]; int score; Trie() { score = 0; fill(node, node+26, (Trie *)0); } void insert(const string &s, int val) { Trie* r = this; for(size_t i=0; i<s.length(); i++) { int c = s[i] - 'a'; if(!r -> node[c]) r -> node[c] = new Trie; r = r -> node[c]; } r -> score = val; } Trie* find(const char &s, Trie* pos) { int c = s - 'a'; if(!pos -> node[c]) return (Trie *)0; return pos -> node[c]; } }; string p[5010]; int score[5010], dp[200010]; signed main() { Trie trie; string s; cin >> s; int M, N; cin >> M; N = s.length(); rep(i,0,M) cin >> p[i]; rep(i,0,M) cin >> score[i]; rep(i,0,M) trie.insert(p[i], score[i]); rep(i,0,N) { Trie* cur = ≜ repq(k,1,200) { if(i+k > N) continue; char target = s[i+k-1]; Trie* next; if(cur != (Trie *)0) { next = trie.find(target, cur); cur = next; chmax(dp[i+k], (cur != (Trie *)0) ? dp[i] + cur -> score : dp[i]); } else chmax(dp[i+k], dp[i]); } } cout << dp[N] << endl; return 0; }
ハマったところ
- 出現する文字は英小文字のみなので Trie のポインタ配列は 26 あれば十分 (ライブラリでは 256 にしていて、それをそのまま使うと MLE した)
- find は前の結果をうまく使わないと遅くなる (いちいち trie 木の根からやるのはアウト)
- 左端を固定するとき、ある部分文字列に対してマッチするものがなければ、その後右端を伸ばしてもマッチするはずはないが、値の伝搬は忘れない (break で抜けると伝搬が起こらないのでアウト)