PythonでTrie（プレフィックス木）を実装する方法をわかりやすく解説

Trie（トライ木・プレフィックス木）は、文字列の集合を効率的に管理するための木構造データ構造です。検索エンジンの補完機能やスペルチェッカーなど、 prefix（接頭辞）を扱う処理で広く活用されています。

この記事では、Pythonを使ってTrieを実装する方法を解説します。実装するのは以下の3つの基本操作です。

• insert(word)：単語をTrieに挿入する
• search(word)：指定した単語が完全一致で登録されているか判定する
• startsWith(prefix)：指定した接頭辞で始まる単語が存在するか判定する

動作イメージ

すべての入力は小文字の英字であると仮定します。例えば、次のようにメソッドを呼び出した場合の出力を見てみましょう。

• Trie trie = new Trie()
• trie.insert("apple")
• trie.search("apple")　// true を返す
• trie.search("app")　　// false を返す（"app" はまだ挿入されていないため）
• trie.startsWith("app") // true を返す（"apple" の接頭辞だから）
• trie.insert("app")
• trie.search("app")　　// true を返す

このように、search() は単語が完全に登録されているかを判定するのに対し、startsWith() は接頭辞として一致する経路があるかだけを判定する点がポイントです。

アルゴリズムの考え方

ここでは、Pythonの辞書（dict）をノード代わりに使うシンプルな手法を採用します。各ノードは「文字 → 子ノードの辞書」という対応を持ち、単語の終端には特別なキー # を置くことで「ここで単語が完結している」ことを表します。

insert メソッドの手順

1. current を child（ルートの辞書）に設定する
2. word の各文字 l について：
   • l が current に存在しなければ、current[l] に新しい辞書を作成する
   • current を current[l] に進める
3. 最後に current['#'] = 1 を設定して単語の終端をマークする

search メソッドの手順

1. current を child に設定する
2. word の各文字 l について：
   • l が current に存在しなければ false を返す
   • current を current[l] に進める
3. current に '#' が存在すれば true、なければ false を返す

startsWith メソッドの手順

1. current を child に設定する
2. prefix の各文字 l について：
   • l が current に存在しなければ false を返す
   • current を current[l] に進める
3. 最後まで辿れれば true を返す（終端マークは不要）

Pythonでの実装コード

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

class Trie(object):
    def __init__(self):
        self.child = {}

    def insert(self, word):
        current = self.child
        for l in word:
            if l not in current:
                current[l] = {}
            current = current[l]
        current['#'] = 1

    def search(self, word):
        current = self.child
        for l in word:
            if l not in current:
                return False
            current = current[l]
        return '#' in current

    def startsWith(self, prefix):
        current = self.child
        for l in prefix:
            if l not in current:
                return False
            current = current[l]
        return True


ob1 = Trie()
ob1.insert("apple")
print(ob1.search("apple"))
print(ob1.search("app"))
print(ob1.startsWith("app"))
ob1.insert("app")
print(ob1.search("app"))

入力

ob1 = Trie()
ob1.insert("apple")
print(ob1.search("apple"))
print(ob1.search("app"))
print(ob1.startsWith("app"))
ob1.insert("app")
print(ob1.search("app"))

出力

True
False
True
True

処理の流れを確認しよう

このコードを実行すると、まず "apple" を挿入します。search("apple") は完全一致する単語が登録済みなので True を返します。一方、search("app") は「a → p → p」という経路自体は存在しますが、その時点で単語が完結していない（# マークがない）ため False となります。

しかし startsWith("app") は接頭辞としての経路があればよいので True を返します。その後 "app" を挿入すると、search("app") も True になることが確認できます。

計算量について

この実装における計算量は以下の通りです。

• insert / search / startsWith：いずれも O(m)（m は操作対象の文字列の長さ）。文字数分だけ辞書を辿るだけで完了します。
• 空間計算量：保存した単語の総文字数に依存しますが、共通の接頭辞はノードを共有するため、個別に文字列を保存するより効率的です。

辞書ベースのTrieはコードが簡潔で理解しやすいのが魅力です。より高度な用途では、子ノードを配列で持つ実装や、ノード数を削減した圧縮Trie（Radix Tree）なども検討するとよいでしょう。