説明

プレフィックス木とも呼ばれる。あるノードの配下の全てのノードには自身に対応する文字列に共通するプレフィックスがあり, ルートには空の文字列が対応している。

計算量

  • 追加 $O(|S|)$
  • クエリ $O(|S|)$

実装例

exist は子供以下に追加された文字列の個数, accept はそのノードにマッチする全ての追加された文字列の番号が格納される。

  • Trie< char_size, margin >(): 文字列の種類数 char_size, 開始文字が margin のトライ木を作成する。
  • add($S$): 文字列 $S$ をトライ木に追加する
  • query($S$, $f$): 文字列 $S$ のプレフィックスに一致する文字列を検索する。一致した文字列ごとに関数 $f$ が呼び出される。
  • size(): ノード数を返す
  • count(): 存在する文字列の個数を返す
template< int char_size >
struct TrieNode {
  int nxt[char_size];

  int exist;
  vector< int > accept;

  TrieNode() : exist(0) {
    memset(nxt, -1, sizeof(nxt));
  }
};

template< int char_size, int margin >
struct Trie {
  using Node = TrieNode< char_size >;

  vector< Node > nodes;
  int root;

  Trie() : root(0) {
    nodes.push_back(Node());
  }

  void update_direct(int node, int id) {
    nodes[node].accept.push_back(id);
  }

  void update_child(int node, int child, int id) {
    ++nodes[node].exist;
  }

  void add(const string &str, int str_index, int node_index, int id) {
    if(str_index == str.size()) {
      update_direct(node_index, id);
    } else {
      const int c = str[str_index] - margin;
      if(nodes[node_index].nxt[c] == -1) {
        nodes[node_index].nxt[c] = (int) nodes.size();
        nodes.push_back(Node());
      }
      add(str, str_index + 1, nodes[node_index].nxt[c], id);
      update_child(node_index, nodes[node_index].nxt[c], id);
    }
  }

  void add(const string &str, int id) {
    add(str, 0, 0, id);
  }

  void add(const string &str) {
    add(str, nodes[0].exist);
  }

  void query(const string &str, const function< void(int) > &f, int str_index, int node_index) {
    for(auto &idx : nodes[node_index].accept) f(idx);
    if(str_index == str.size()) {
      return;
    } else {
      const int c = str[str_index] - margin;
      if(nodes[node_index].nxt[c] == -1) return;
      query(str, f, str_index + 1, nodes[node_index].nxt[c]);
    }
  }

  void query(const string &str, const function< void(int) > &f) {
    query(str, f, 0, 0);
  }

  int count() const {
    return (nodes[0].exist);
  }

  int size() const {
    return ((int) nodes.size());
  }
};

検証

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

int main() {
  Trie< 26, 'a' > trie;
  string S, P[5000];
  int M, W[5000];

  cin >> S;
  cin >> M;
  for(int i = 0; i < M; i++) {
    cin >> P[i];
    trie.add(P[i]);
  }
  for(int i = 0; i < M; i++) {
    cin >> W[i];
  }
  vector< int > dp(S.size() + 1);
  for(int i = 0; i < S.size(); i++) {
    auto update = [&](int idx) {
      dp[i + P[idx].size()] = max(dp[i + P[idx].size()], dp[i] + W[idx]);
    };
    trie.query(S, update, i, 0);
    dp[i + 1] = max(dp[i + 1], dp[i]);
  }
  cout << dp.back() << endl;
}