トライ木(Trie木) の解説と実装【接頭辞(prefix) を利用したデータ構造】

2020年3月27日データ構造配列,Trie木,文字列,prefix

Trie木は、効率的な検索(retrieval)のために使われるデータ構造です。文字列などの先頭部分(接頭辞: prefix)の共通部分を共有して保存することで、\(O(M)\) での検索を可能にします。(文字列の長さを\(M\) としました。)

以下は、「fire・firearm・firework・fireman・algo」をトライ木に保存した時の状態を表しています。

トライ木の仕組み

基本は木

複数の頂点から構成された順序付きの木になっています。効率的な検索のために

  • 1つの頂点が一つの文字を表す
  • 文字列の先頭部分(prefix)が共通する場合は、なるべく無駄が無いように共有する

というようになっています。

頂点

トライ木は、一つの頂点(ノード)が一文字を表現した木となっています。必要に応じて様々な情報を頂点が持たせると実装しやすいです。基本としては、

  • その頂点の子は何か
  • その頂点で末端となる文字列はあるか
  • その頂点はどんな文字列を表現しているか

などが必要になるでしょう。

トライ木の機能

基本的な機能としては以下の2つです。

  • 挿入(insert):文字列をトライ木に保存する
  • 検索(search):文字列がトライ木に保存されているか確認する

文字列の挿入

挿入時は、文字列を一文字ずつ確認して、Trie 木の頂点を1つずつたどっていきます。もしまだ頂点が存在しなければ、新しく作成することになります。

「firework」 の挿入前
「firework」の挿入後

文字列の検索

検索時も、文字列を一文字ずつ確認して、Trie 木の頂点を1つずつたどっていきます。もし頂点が存在しなければ、その文字列は Trie 木に挿入されていなかったということになります。

最後の頂点にたどり着いたら、その頂点で終了する文字列があるかを確認する必要があります。例えば、"firework" を挿入したなら、true を返すのは “firework" を検索した時で、"firewo" などを検索した時には false を返してもらうためです。

「firework」で一致
「firewo」では不一致

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;
}

練習問題