天下一プログラマーコンテスト 2016 C: たんごたくさん

初めて 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 = &trie;
        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 で抜けると伝搬が起こらないのでアウト)