マジで実装が簡単なTrieの話

皆さん、突然ですが Trie(データ構造) って実装したいですか? できればしたくないですよね...
コンテスト中に Trie を使う問題に出会ったり、 Trie を用いた実装をしたりすると言っても、 Trie を使ってやりたいことの種類にもいろいろあって使いやすいライブラリを作るのもけっこう難しいと思います。
そこで、この記事では 最小の労力で Trie の本質を記述し、応用もかなり利きやすそうな実装法 を書いてみたいと思います。
なお、この記事は Trie をソラで書くシチュエーションを想定しています。もしもあなたが強力で使い慣れたライブラリを持っているならそれを使うべきです。
今回は、 たんごたくさん(天下一プログラマーコンテスト2016本選-C) を題材として、 C++ で実装を進めていきたいと思います。

Step0. 今回扱う問題の解き方

本題に入る前に、今回扱う問題の解き方が理解できなければ話が始まらないので簡単に解き方を説明しておきます。
ひとまず、問題概要をここにも載せます。

文字列 $S$ と、要素数 $M$ の単語の集合 $P$ が与えられる。
$P$ の各要素には重み $W(>0)$ が定義されている。
文字列 $S$ から $P$ に含まれる(連続)部分列を重なり合わないように取り出す。
うまく取り出し方を決めて、取り出した単語の重みの総和を最大化せよ。
なお、同じ単語を複数回取り出した場合、それらの単語は別々に数える。

先に今回扱う問題の解き方を書いてしまいましょう。
$dp[i]=$ (文字列 $S$ の前 $i$ 文字目までを消費した時に達成可能な最大の重みの総和) で解くことができます。単語の長さが高々 $200$ なので、全ての整数 $1 \le i \le |S|$ について多くとも 「$S_i$ を取り出した場合」「 $S_iS_{i+1}$ を取り出した場合」 $\dots$ 「$S_iS_{i+1} \dots S_{i+199}$ を取り出した場合」に対して遷移を行えばよいです。
要するに、全ての $i$ に対して、 $S_i$ を先頭とした連続部分列について、その部分列を単語の集合 $P$ の Trie 上で辿っていって遷移を行う感じになります。

今回は解法にはこれ以上深入りしません。他にもこの問題の解説をしている記事があるので解法を理解した上で先に進んでください。

Step1. Trie で扱う情報を定義する

まず、今回作りたい Trie の各頂点にどのような情報を持たせたいかを決めます。今回は以下の情報を持たせることになります。

• その頂点の子の情報
• その頂点が表す文字列の重み

今回はこの $2$ つです。構造体で書くと以下の通りです。

typedef struct{
  int child[26]; // その頂点の子の情報
  long long weight; // その頂点が表す文字列の重み
}node;

今回はこの構造体に「その頂点の親の情報」を持たせていません。「 Trie を上る」という操作が必要ない時はその頂点の親の情報は必要ないので実装を軽くすることができます。なお、親の情報が必要である場合も、この構造体にその要素を追加するだけで情報の定義は完了です。

なお、子の情報も必ずしも配列で持つ必要はなく、必要や用途に応じて他のデータ構造を使うこともできます。

Step2. 実際に Trie を構成する

まず、初期化用のノードを作ります。この初期化用のノードは子を持たず ($-1$ を子を持たない状態とする)、頂点に乗った他の情報も適切な初期値を設定しておきます。(今回の問題では、 (重み) $=0$ が適切な初期値です)

// 初期化用のノードを作る
node init;
for(int i=0;i<26;i++){init.child[i]=-1;}
init.weight=0;

では、実際に単語の集合から Trie を構成していきましょう。この部分は実装を見てもらった方が早いと思うので、実装を示します。

vector<node> trie={init};
for(int word=0;word<m;word++){
  int vertex=0; // 頂点 0 は根とする
  for(int i=0;i<p[word].size();i++){
    if(trie[vertex].child[p[word][i]-'a']==-1){ // 今から行きたい子が無い場合
      trie[vertex].child[p[word][i]-'a']=trie.size(); // 次に作られる頂点の添え字は trie.size() なので、その頂点を子として登録する
      trie.push_back(init); // 新たに頂点を作成する
    }
    vertex=trie[vertex].child[p[word][i]-'a']; // 子に移動する
  }
  trie[vertex].weight=w[word]; // 重みを登録する
}

まず、 Trie の根ノードを $0$ とします。その後、各単語ごとに根ノードから下っていくのですが、ある単語を Trie に追加したい場合は、その単語をたどる途中で「子が無い」という事態に遭遇するごとに新たに子を作ることになります。もし親の情報が必要な Trie を作りたい場合は、ここで作った子に対して親の情報をこのタイミングで登録することになります。

ある単語を Trie に追加したくない場合(つまり、単語を単に検索したいだけの場合)もほぼ同じ実装でよく、実装の差異は「子が無いという事態に遭遇した場合、検索結果が(これ以上)ないものとして Trie をたどるのを打ち切って良い」という点だけです。詳しくは以下の「最終的な AC コード」を参考にしてください。

最終的な AC コード

最終的な AC コードは以下です。 提出

#include<bits/stdc++.h>

using namespace std;

typedef struct{
  int child[26]; // その頂点の子の情報
  long long weight; // その頂点が表す文字列の重み
}node;

int main(){
  string s;
  int m;
  cin >> s >> m;
  vector<string> p(m);
  vector<long long> w(m);
  for(int i=0;i<m;i++){cin >> p[i];}
  for(int i=0;i<m;i++){cin >> w[i];}

  // 初期化用のノードを作る
  node init;
  for(int i=0;i<26;i++){init.child[i]=-1;}
  init.weight=0;

  vector<node> trie={init};
  for(int word=0;word<m;word++){
    int vertex=0; // 頂点 0 は根とする
    for(int i=0;i<p[word].size();i++){
      if(trie[vertex].child[p[word][i]-'a']==-1){ // 今から行きたい子が無い場合
        trie[vertex].child[p[word][i]-'a']=trie.size(); // 次に作られる頂点の添え字は trie.size() なので、その頂点を子として登録する
        trie.push_back(init); // 新たに頂点を作成する
      }
      vertex=trie[vertex].child[p[word][i]-'a']; // 子に移動する
    }
    trie[vertex].weight=w[word]; // 重みを登録する
  }

  vector<long long> dp(s.size()+1,0);
  for(int i=0;i<s.size();i++){ // s[i] を先頭とした連続部分列の遷移を行う
    int vertex=0; // 頂点 0 は根とする
    for(int j=i;j<s.size();j++){
      if(trie[vertex].child[s[j]-'a']==-1){break;} // 今から行きたい子がないなら、そこから先を見ても正の重みの単語は得られないので、遷移終了
      vertex=trie[vertex].child[s[j]-'a']; // Trie を下る
      dp[j+1]=max(dp[i]+trie[vertex].weight,dp[j+1]);
    }
    dp[i+1]=max(dp[i],dp[i+1]); // s[i] を使わないことにする遷移
  }

  long long res=0;
  for(auto &nx : dp){res=max(res,nx);}
  cout << res << '\n';
  return 0;
}

この実装の強み

最後に、この実装の強みを軽く説明しておくと、

• とにかくコード量が少ないので、ライブラリが持ち込めないコンテストや状況に対して非常に強力！
  • 実装に際して覚えておかなければならないポイントもほぼないと思います
• 単純な操作を少ない回数行うだけで Trie の本質部分の実装が完了する！
• 行っている操作が簡単なので、高速！
• 頂点を管理する構造体をいじってやれば応用幅も広い！

あたりがこの実装の強みです。
Trie の実装と言えば 重い / 面倒 というイメージを受けてしまっていたかもしれませんが、この実装法だと相当楽に見えてくるのではないでしょうか。
では、皆さんよき Trie ライフを！