ソラで書くTrie

Trieは文字列集合をラベル付き木構造で表現したデータ構造で，以下のような特徴を持ちます．

1. 文字列集合内の任意の接頭辞とノードが1対1に対応する．
2. 文字列を入力としてTrie上を探索することで，文字列集合内のキーを効率的に検索できる．
3. 文字列集合の共通接頭辞が併合され，コンパクトな表現ができる．

1の特徴は文字列集合に対する様々なアルゴリズムへ応用され，
2の特徴から様々な文字列索引技術の基礎として利用され，
3の特徴から大規模データの圧縮にも利用されます．

特にTrieを応用してアルゴリズムを設計する場合，Trieのアルゴリズムをよく理解しスムーズに実装できることが重要です．
そこでまず本記事では，実装の柔軟性を考慮し，簡単なデータ構造の組み合わせによるTrieの実装をレクチャーします．

Trieの実装

Trieで重要な基本操作は以下の2つです．

• $\texttt{go}(v, c) \rightarrow V$: ノード$v$からラベル$c$で遷移するノード$t$があれば$t$を，なければ−1を返す
• $\texttt{add_node}(v, c) \rightarrow V$: ノード$v$からラベル$c$で遷移するノード$t$を追加し，$t$を返す

$\texttt{go}, \texttt{add_node}$さえ実装すれば，Trieを用いたアルゴリズムは基本的に実装できます．
$\texttt{go}, \texttt{add_node}$の実装は後で解説するとして，まずTrieを用いた簡単なアルゴリズムを紹介します．

文字列を入力とした探索

文字列$s$を入力とし，$s$による探索で到達するノード$v$があれば$v$を，なければ−1を返すことにします．
これは，$s$の先頭から一文字ずつ取り出して遷移ラベル$c$とし,ノード$u=0(\text{ルートノード})$から順番に$\texttt{go}(u,c)$で遷移を繰り返せばよいです．

int travel(const std::string& s) {
  int v = 0;
  for (char c : s) {
    int t = go(v, c);
    if (t != -1)
      v = t;
    else
      return -1;
  }
  return v;
}

文字列の追加

文字列$s$をTrieに追加します．
これは，まず$s$による探索を遷移に失敗するまで繰り返します．
その後，残りの文字列から一文字ずつ取り出して遷移先となるノードを順に追加していきます．

void insert(const std::string& s) {
  int v = 0;
  int i = 0;
  for (; i < (int) s.size(); i++) {
    int t = go(v, s[i]);
    if (t != -1)
      v = t;
    else
      break;
  }
  for (; i < (int) s.size(); i++) {
    v = add_node(v, s[i]);
  }
}

データ構造

もっとも単純なデータ構造のフレームワークは，
遷移文字をキーとしてノード番号を保存する索引構造Indexを与え，
Indexをノードごとに管理するのものでしょう．

using Index = // 任意の索引データ構造
std::vector<Index> trie(1); // ルートノード=0で初期化

このフレームワークでIndexを任意のデータ構造にするだけで，以下の要件を満たすTrieを簡単に設計できます．

|Index|メモリサイズ|検索|全探索(列挙)|
|-|-|-|-|-|
|Table|×|◎|×|
|Array|○|×|◎|
|HashTable|○|○|○¹|

Index = Table

Indexを固定長配列で与え，遷移先の状態番号を保存します．
すなわち，ラベル付きグラフ表現としてもっとも単純な方法である遷移テーブルとしてTrieを表現します．

この利点は，任意のラベルによる遷移が，ラベル$c$を添え字とした配列により定数時間で行える点です．
欠点は，ノードごとにアルファベット集合のサイズ分の領域が必要であり，メモリを大量に消費するため，扱うデータサイズに注意する必要があります．

constexpr int kAlphabetSize = 256 // アルファベット集合のサイズ
using Index = std::array<int, kAlphabetSize>;
std::vector<Index> trie(1);

int go(int v, char c) {
  int next_v = trie[v][c];
  return next_v != 0 ? next_v : -1;
}

int add_node(int v, char c) {
  int next_v = (int) trie.size();
  trie.emplace_back();
  trie[v][c] = next_v;
  return next_v;
}

Index = Array

Indexを動的配列で与え，（遷移ラベル，遷移先の状態番号）の組を保存します．

利点は，要素数がノード数で抑えられるためTableよりコンパクトになります．
また，トライの全ノードの探索や保存された文字列の列挙を行う場合には，要素を順に見るだけでよくTableより効率的になります．
欠点は，任意ラベルによる探索には，ノード$v$のIndexの要素を順に見る必要があり，検索が非効率になります．

using Index = std::vector<pair<char, int>>;
std::vector<Index> trie(1);

int go(int v, char c) {
  for (auto [cur_c, next_v] : trie[v]) {
    if (cur_c != c) 
      continue;
    return next_v != 0 ? next_v : -1;
  }
  return -1;
}

int add_node(int v, char c) {
  int next_v = (int) trie.size();
  trie.emplace_back();
  trie[v].emplace_back(c, next_v);
  return next_v;
}

Index = HashTable

Indexをハッシュテーブルで与え，遷移ラベルをキーとして遷移先の状態番号を保存します．

TableやArrayと違い，空間効率と検索効率を兼ね備えることができます．
ただし，検索速度はTableの方が定数倍高速です．

using Index = std::unordered_map<char, int>;
std::vector<Index> trie(1);

int go(int v, char c) {
  auto node_it = trie[v].find(c);
  return node_it != trie[v].end() ? node_it->second : -1;
}

int add_node(int v, char c) {
  int next_v = (int) trie.size();
  trie.emplace_back();
  trie[v][c] = next_v;
  return next_v;
}

まとめ

ノードごとにデータ構造を与える単純なフレームワークを用いたTrieの実装を紹介しました．
Tableにはメモリサイズの,Arrayには検索効率の欠点があり，注意深く用いる必要がありますが，欠点を無視できる状況ならばこれらの簡単な構造は有力な選択肢になります．

競技プログラミングではTrieを応用することを想定した問題が度々出題されるため，ソラで書ける能力が有利に働くこともあるでしょう．

• 東京海上日動 プログラミングコンテスト2021（AtCoder Regular Contest 122) D - XOR Game
  • 回答例

[編集歓迎]

1. std::unordered_mapのように要素がリスト構造で連結されている構造を仮定 ↩