トライ木(Trie木) の解説と実装【接頭辞(prefix) を利用したデータ構造】
Trie木は、効率的な検索(retrieval)のために使われるデータ構造です。文字列などの先頭部分(接頭辞: prefix)の共通部分を共有して保存することで、\(O(M)\) での検索を可能にします。（文字列の長さを\(M\) としました。）
以下は、「fire・firearm・firework・fireman・algo」をトライ木に保存した時の状態を表しています。
トライ木の仕組み
基本は木
複数の頂点から構成された順序付きの木になっています。効率的な検索のために
- 1つの頂点が一つの文字を表す
- 文字列の先頭部分(prefix)が共通する場合は、なるべく無駄が無いように共有する
というようになっています。
頂点
トライ木は、一つの頂点(ノード)が一文字を表現した木となっています。必要に応じて様々な情報を頂点が持たせると実装しやすいです。基本としては、
- その頂点の子は何か
- その頂点で末端となる文字列はあるか
- その頂点はどんな文字列を表現しているか
などが必要になるでしょう。
トライ木の機能
基本的な機能としては以下の2つです。
- 挿入(insert)：文字列をトライ木に保存する
- 検索(search)：文字列がトライ木に保存されているか確認する
文字列の挿入
挿入時は、文字列を一文字ずつ確認して、Trie 木の頂点を1つずつたどっていきます。もしまだ頂点が存在しなければ、新しく作成することになります。
文字列の検索
検索時も、文字列を一文字ずつ確認して、Trie 木の頂点を1つずつたどっていきます。もし頂点が存在しなければ、その文字列は Trie 木に挿入されていなかったということになります。
最後の頂点にたどり着いたら、その頂点で終了する文字列があるかを確認する必要があります。例えば、"firework" を挿入したなら、true を返すのは “firework" を検索した時で、"firewo" などを検索した時には false を返してもらうためです。
C++での実装例
ベース
実装方法は様々なものがありますが、頂点の集合を vector などに保存しておくことにして実装していくことにします。C++では、以下のように構造体などで実装すると変更がしやすいです。
template <int char_size, int base>
struct Trie {
struct Node { // 頂点を
vector<int> next; // 子の頂点番号を格納。存在しなければ-1
vector<int> accept; // 末端がこの頂点になる文字列の str_id を保存
int c; // base からの間隔をint型で表現したもの
int common; // いくつの文字列がこの頂点を共有しているか
Node(int c_) : c(c_), common(0) {
next.assign(char_size, -1);
}
};
vector<Node> nodes; // trie 木本体
int root; // 根
Trie() : root(0) {
// 初期化。はじめは根しか無い
nodes.push_back(Node(root));
}
/*
その他のメソッド（insert や search） などを実装する
*/
}
※ 頂点は子の頂点番号を保持することにしていますが、子の頂点への参照を保持させるような実装もあります。その場合は、少し実装が複雑になりますが、頂点の集合を vector で保持しておく必要などはありません。
※ 文字の種類は char_size に限定されているとして、その 0 番目に当たるものを base としています。例えば、Trie<26, 'a’> であれば、アルファベット小文字26種類からなる文字列のみを考えることになります。このようにすることで、「’a’ を数字の 0」「’c’を数字の2」などと簡単に対応させることができます。
単語の挿入
一文字ずつ頂点が存在するか確認していきましょう。存在しない場合には、頂点を新しく追加する必要があります。
最後の頂点には、単語の終端だと分かるように単語の番号を加えています。
void insert(const string &word, int word_id) { // 単語:str と 単語の番号: str_id
int node_id = 0;
for (int i = 0; i < (int)word.size(); i++) {
int c = (int)(word[i] - base);
int &next_id = nodes[node_id].next[c];
if (next_id == -1) { // 次の頂点が存在しなければ追加
next_id = (int)nodes.size();
nodes.push_back(Node(c));
}
++nodes[node_id].common;
node_id = next_id;
}
++nodes[node_id].common;
nodes[node_id].accept.push_back(word_id); // 単語の終端なので、頂点に番号を入れておく
}
void insert(const string &word) {
insert(word, nodes[0].common);
}
単語の検索
単語を1文字ずつ確認していき、次の頂点が存在しなければ終了します。
全ての頂点があったとしても、最後の頂点が受理状態か確認する必要があります。
// 単語とprefixの検索
bool search(const string &word, bool prefix = false) {
int node_id = 0;
for (int i = 0; i < (int)word.size(); i++) {
int c = (int)(word[i] - base);
int &next_id = nodes[node_id].next[c];
if (next_id == -1) { // 次の頂点が存在しなければ終了
return false;
}
node_id = next_id;
}
return nodes[node_id].accept.size() > 0; // 最後の頂点が受理状態か確認
}
その他の機能
いくつかあると便利な機能を加えておきます。プログラミングコンテストなどでは、これらの実装例そのままで解ける問題は少ないですが、トライ木を利用するための元になります。
指定した prefix を持つ文字列があるか検索
search のコードを一部変更することで、単語の検索だけではなく、指定した接頭辞(prefix) を持つ単語が存在するか判定することができます。終端の頂点が受理状態であろうと無かろうと true を返すようにすれば良いです。
例えば、「fireman」を挿入していれば、「firem」を接頭辞として持つことになります。
// 単語とprefixの検索
bool search(const string &word, bool prefix = false) {
int node_id = 0;
for (int i = 0; i < (int)word.size(); i++) {
int c = (int)(word[i] - base);
int &next_id = nodes[node_id].next[c];
if (next_id == -1) { // 次の頂点が存在しなければ終了
return false;
}
node_id = next_id;
}
return (prefix) ? true : nodes[node_id].accept.size() > 0; // 最後の頂点が受理状態か確認
}
// prefix を持つ単語が存在するかの検索
bool start_with(const string &prefix) {
return search(prefix, true);
}
挿入した単語の総数
全ての頂点に、その頂点を共有している単語の数を記録してあるので、根(root) の頂点について確認すれば、挿入した単語の総数を確認することができます。
// 挿入した単語の数
int count() const {
return (nodes[0].common);
}
頂点の総数
prefix を共有することで、いくつか頂点を節約できているはずです。この頂点は vector で管理していたので、簡単に求めることができます。
// Trie木のノード数
int size() const {
return ((int)nodes.size());
}
C++ での実装まとめ
上述の全ての機能を盛り込んだ実装のまとめです。
#include <bits/stdc++.h>
using namespace std;
/* Trie 木： 文字の種類(char_size)、int型で0に対応する文字(base)
insert(word): 単語 word を Trie 木に挿入する
search(word): 単語 word が Trie 木にあるか判定する
start_with(prefix): prefix が一致する単語が Trie 木にあるか判定する
count(): 挿入した単語の数を返す
size(): Trie 木の頂点数を返す
計算量：insert, search ともに O(M)（Mは単語の長さ）
*/
template <int char_size, int base>
struct Trie {
struct Node { // 頂点を表す構造体
vector<int> next; // 子の頂点番号を格納。存在しなければ-1
vector<int> accept; // 末端がこの頂点になる単語の word_id を保存
int c; // base からの間隔をint型で表現したもの
int common; // いくつの単語がこの頂点を共有しているか
Node(int c_) : c(c_), common(0) {
next.assign(char_size, -1);
}
};
vector<Node> nodes; // trie 木本体
int root;
Trie() : root(0) {
nodes.push_back(Node(root));
}
// 単語の挿入
void insert(const string &word, int word_id) {
int node_id = 0;
for (int i = 0; i < (int)word.size(); i++) {
int c = (int)(word[i] - base);
int &next_id = nodes[node_id].next[c];
if (next_id == -1) { // 次の頂点が存在しなければ追加
next_id = (int)nodes.size();
nodes.push_back(Node(c));
}
++nodes[node_id].common;
node_id = next_id;
}
++nodes[node_id].common;
nodes[node_id].accept.push_back(word_id);
}
void insert(const string &word) {
insert(word, nodes[0].common);
}
// 単語とprefixの検索
bool search(const string &word, bool prefix = false) {
int node_id = 0;
for (int i = 0; i < (int)word.size(); i++) {
int c = (int)(word[i] - base);
int &next_id = nodes[node_id].next[c];
if (next_id == -1) { // 次の頂点が存在しなければ終了
return false;
}
node_id = next_id;
}
return (prefix) ? true : nodes[node_id].accept.size() > 0;
}
// prefix を持つ単語が存在するかの検索
bool start_with(const string &prefix) {
return search(prefix, true);
}
// 挿入した単語の数
int count() const {
return (nodes[0].common);
}
// Trie木のノード数
int size() const {
return ((int)nodes.size());
}
};
int main() {
Trie<26, 'A'> trie;
trie.insert("FIRE");
cout << trie.search("FIRE") << endl; // 1 を返す
cout << trie.search("FI") << endl; // 0 を返す
trie.insert("FIREMAN");
trie.insert("FIREARM");
trie.insert("FIREWORK");
trie.insert("ALGO");
return 0;
}
練習問題
- Implement Trie (Prefix Tree) – LeetCode：単純な Trie 木を実装するだけの問題です
- Google Kick Start 2020 Round A Bundling ：Trie 木を上手く利用すると簡単に解くことができます