Trie（トライ木）で作るオートコンプリート機能｜Pythonでの実装方法を解説

Trie（トライ木）によるオートコンプリート機能とは？

Trie（トライ木・接頭辞木）は、文字列の集合を効率的に管理できるデータ構造の一つです。このTrieを活用すると、ユーザーが文字を入力するたびに一致する候補を即座に表示するオートコンプリート（自動補完）機能を実現できます。

例えば、Trieに 「xyz」「xyzzzz」「xyabad」… といった単語が登録されている状態で、ユーザーが xy と入力すると、「xyz」「xyzzzz」「xyabad」など、入力した接頭辞で始まる単語がすべて候補として表示されます。検索エンジンのサジェストやエディタのコード補完など、身近な多くの機能がこの仕組みを応用しています。

オートコンプリートを実現する手順

1. 接頭辞の検索：標準的なTrie探索アルゴリズムで、入力文字列をルートから順にたどる。
2. 見つからない場合：該当するパスが存在しなければ -1 を返し、一致なしと判定する。
3. 完全一致の場合：入力文字列が単語の終端であり、かつ子ノードを持たない場合は、その文字列だけを出力する。
4. 子ノードがない場合：一致したノードにそれ以上の子ノードがなければ、そこで処理を終了する。
5. 候補の一覧表示：上記以外の場合は、一致したノード以下を再帰的にたどり、すべての候補文字列を出力する。

Pythonでの実装コード

それでは、実際のコードを見ていきましょう。

# Trieのノードを表すクラス
class TrieNode():
    def __init__(self):
        # 子ノードを格納する辞書
        self.trie_node = {}
        # 単語の終端かどうかを示すフラグ
        self.last_node = False


class Trie():
    def __init__(self):
        # ルートノードを初期化
        self.root = TrieNode()
        # 補完候補を保存するリスト
        self.words = []

    def create_trie(self, keys):
        # キーのリストからTrieを構築する
        for key in keys:
            self.insert_node(key)

    def insert_node(self, key):
        node = self.root
        for obj in list(key):
            if not node.trie_node.get(obj):
                # 新しいノードを作成
                node.trie_node[obj] = TrieNode()
            node = node.trie_node[obj]
        # 葉ノードであることを記録
        node.last_node = True

    def search(self, key):
        # キーを検索する
        node = self.root
        is_found = True
        for obj in list(key):
            if not node.trie_node.get(obj):
                is_found = False
                break
            node = node.trie_node[obj]
        return node and node.last_node and is_found

    def matches(self, node, word):
        # 単語の終端なら結果リストに追加
        if node.last_node:
            self.words.append(word)
        # 子ノードを再帰的にたどる
        for obj, n in node.trie_node.items():
            self.matches(n, word + obj)

    def show_auto_completion(self, key):
        node = self.root
        is_found = False
        temp = ''
        for obj in list(key):
            # 1文字ずつ照合する
            if not node.trie_node.get(obj):
                is_found = True
                break
            temp += obj
            node = node.trie_node[obj]
        if is_found:
            # 一致するパスが存在しない
            return 0
        elif node.last_node and not node.trie_node:
            # 完全一致し、それ以上の候補がない
            return -1
        # 一致するすべての候補を収集して表示
        self.matches(node, temp)
        for string in self.words:
            print(string)
        return 1


# Trieに登録するデータ
strings = ["xyz", "xyzzzz", "xyabad", "xyyy",
           "abc", "abbccc", "xyx", "xyxer", "a"]

# オートコンプリート対象の入力文字列
string = "xy"

status = ["Not found", "Found"]

# Trieクラスのインスタンスを生成
trie = Trie()

# 文字列リストからTrieを構築
trie.create_trie(strings)

# 入力文字列に対する補完候補を取得
result = trie.show_auto_completion(string)

if result == -1 or result == 0:
    print("No matches")

実行結果

上記のコードを実行すると、次のような出力が得られます。

xyz
xyzzzz
xyabad
xyyy
xyx
xyxer

入力 xy で始まる6つの単語がすべて表示されています。一方、「abc」「abbccc」「a」のように接頭辞が一致しない単語は候補に含まれないことが確認できます。

コードのポイント解説

1. TrieNodeクラス

各ノードは、子ノードを保持する辞書 trie_node と、そのノードが単語の終端であることを示すフラグ last_node を持ちます。これにより、「xy」というパスが途中までしか存在しないケースと、ひとつの完成した単語になっているケースを正確に区別できます。

2. insert_nodeメソッド

挿入したい文字列を1文字ずつ取り出し、対応する子ノードが存在しなければ新しく作成しながら進みます。最後に到達したノードの last_node を True にすることで、そこが単語の終わりであるとマークします。

3. matchesメソッド（再帰探索）

オートコンプリートの中核となる部分です。last_node が True のノードに到達するたびに、そこまで組み立てた文字列を結果リストに追加し、さらに子ノードがあれば再帰的に探索を続けます。これにより、指定した接頭辞から始まるすべての単語を漏れなく収集できます。

4. show_auto_completionメソッド

ユーザーの入力を1文字ずつTrie上で照合し、パスが途切れたら 0（一致なし）、完全一致して子ノードも存在しなければ -1 を返します。それ以外の場合は matches メソッドで候補を収集し、画面に出力して 1 を返します。

まとめ

Trieは共通の接頭辞を持つ単語を木構造でまとめて管理できるため、前方一致検索が必要なオートコンプリートの実装に非常に適しています。計算量は入力文字数と候補の総文字数に依存し、単語数が増えてもリストやハッシュテーブルの線形探索より効率的に動作します。検索サジェスト、IME変換、コード補完など幅広い場面で応用できるので、ぜひ自分のプロジェクトでも試してみてください。