[データ構造] TRIE木

Python では、連想配列として辞書型が組み込まれているので、そちらを通常は利用します。 5.5.辞書型 (dictionary) マッピング型 ---dict ここでは辞書は使わず、自分の理解のため簡単な連想配列をスクラッチから実装して...

TRIE木

トライ木: trie)やプレフィックス木: prefix tree)とは、順序付きの一種。あるノードの配下の全ノードは、自身に対応する文字列に共通するプレフィックス(接頭部)があり、ルート(根)には空の文字列が対応している。値は一般に全ノードに対応して存在するわけではなく、末端ノードや一部の中間ノードだけがキーに対応した値を格納している。2分探索木と異なり、各ノードに個々のキーが格納されるのではなく、木構造上のノードの位置とキーが対応している。

出典: フリー百科事典『ウィキペディア(Wikipedia)』

「文字列の検索」のために使われる木構造で、2分探索木やハッシュテーブルの代替として使うことができます。

トライ木では、各ノードを文字列のそこまでの “prefix” に対応させることで文字列を表現します。

2分探索木より全般的に高速に動作します。

ハッシュテーブルと比較すると検索が速く、またキー順によるソートを行うことができます。

しかし、単純な実装ではノードに使う文字全てを保持するので、メモリ効率がとても悪くなります。

TRIE木の単純な実装

Python で単純な TRIE木を実装してみます。

TRIE木の単純な実装では、ルートを空ノードにしてそれぞれのノードから必要なだけエッジを伸ばした木構造(例えばアルファベットをキーに使う場合は26のエッジを伸ばす)を、配列を使って構成するだけです。

より効率的にするためには、配列ではなく、ダブルアレイ、LOUDS、三分探索木といったデータ構造を使うそうです。

【前編】ダブルアレイで高速でコンパクトな辞書を実装する〜準備〜

簡潔データ構造 LOUDS の解説(全12回、練習問題付き)

エッジ数

エッジの数を定数として保持します。

ここでは、小文字のアルファベットのみがキーになると想定して、26文字とします。

ALPHABET_SIZE = 26

ノード

それぞれのノードは、エッジ数分の子を配列で管理し、-1で初期化します。

また、そこで単語が終わったのかどうかを管理する変数、そこまでの文字列でprefix として該当する総数を保持する変数を持ちます。

ALPHABET_SIZE = 26

class Node(object):
    def __init__(self, char):
        self.char = char
        self.children = [-1] * ALPHABET_SIZE
        self.word_finished = False
        self.matched_prefix_cnt = 0

Trie クラス

ルートは空文字で初期化します。

また、小文字アルファベットを配列で管理するためのヘルパー関数 _convert_char_to_num を定義します。

ALPHABET_SIZE = 26

class Node(object):
    def __init__(self, char):
        self.char = char
        self.children = [-1] * ALPHABET_SIZE
        self.word_finished = False
        self.matched_prefix_cnt = 0


class Trie(object):
    def __init__(self):
        self.root = Node(' ')

    def _convert_char_to_num(self, char):
        return ord(char) - ord('a')

単語の追加

木構造にそこまでの prefix が既に含まれるかどうか、ループで単語を一文字ずつ確認し、含まれない場合は文字を追加します。

単語が終わったら、そのノードの終了フラグをTrueにして、その単語の存在を記憶します。

ALPHABET_SIZE = 26

class Node(object):
    def __init__(self, char):
        self.char = char
        self.children = [-1] * ALPHABET_SIZE
        self.word_finished = False
        self.matched_prefix_cnt = 0


class Trie(object):
    def __init__(self):
        self.root = Node(' ')

    def _convert_char_to_num(self, char):
        return ord(char) - ord('a')

    def insert(self, word):
        current = self.root
        for char in word:
            char_index = self._convert_char_to_num(char)
            if current.children[char_index] != -1:
                current = current.children[char_index]
                current.matched_prefix_cnt += 1
            else:
                new_node = Node(char)
                current.children[char_index] = new_node
                current = new_node
                current.matched_prefix_cnt += 1

        current.word_finished = True

単語の検索

ループを使って、一文字ずつ単語が木構造に当てはまるかどうかを確認します。

ALPHABET_SIZE = 26

class Node(object):
    def __init__(self, char):
        self.char = char
        self.children = [-1] * ALPHABET_SIZE
        self.word_finished = False
        self.matched_prefix_cnt = 0


class Trie(object):
    def __init__(self):
        self.root = Node(' ')

    def _convert_char_to_num(self, char):
        return ord(char) - ord('a')

    def insert(self, word):
        current = self.root
        for char in word:
            char_index = self._convert_char_to_num(char)
            if current.children[char_index] != -1:
                current = current.children[char_index]
                current.matched_prefix_cnt += 1
            else:
                new_node = Node(char)
                current.children[char_index] = new_node
                current = new_node
                current.matched_prefix_cnt += 1

        current.word_finished = True

    def search(self, word):
        if not self.root.children:
            return False, 0
        
        current = self.root
        for char in word:
            char_index = self._convert_char_to_num(char)
            if current.children[char_index] != -1:
                current = current.children[char_index]
            else:
                return False, 0
        if current.word_finished:
            return True, current.matched_prefix_cnt
        return False, current.matched_prefix_cnt

テスト

想定通りに動くか、下のような単語を入力して確認します。

‘A’ は ‘a’ にします。

ALPHABET_SIZE = 26

class Node(object):
    def __init__(self, char):
        self.char = char
        self.children = [-1] * ALPHABET_SIZE
        self.word_finished = False
        self.matched_prefix_cnt = 0


class Trie(object):
    def __init__(self):
        self.root = Node(' ')

    def _convert_char_to_num(self, char):
        return ord(char) - ord('a')

    def insert(self, word):
        current = self.root
        for char in word:
            char_index = self._convert_char_to_num(char)
            if current.children[char_index] != -1:
                current = current.children[char_index]
                current.matched_prefix_cnt += 1
            else:
                new_node = Node(char)
                current.children[char_index] = new_node
                current = new_node
                current.matched_prefix_cnt += 1

        current.word_finished = True

    def search(self, word):
        if not self.root.children:
            return False, 0
        
        current = self.root
        for char in word:
            char_index = self._convert_char_to_num(char)
            if current.children[char_index] != -1:
                current = current.children[char_index]
            else:
                return False, 0
        if current.word_finished:
            return True, current.matched_prefix_cnt
        return False, current.matched_prefix_cnt

if __name__ == '__main__':
    trie = Trie()
    trie.insert('a')
    trie.insert('to')
    trie.insert('tea')
    trie.insert('ted')
    trie.insert('ten')
    trie.insert('i')
    trie.insert('in')
    trie.insert('inn')

    # (True, 1)
    print(trie.search('a'))
    # (False, 4)
    print(trie.search('t'))
    # (True, 1)
    print(trie.search('ted'))
    # (False, 0)
    print(trie.search('tel'))
    # (True, 2)
    print(trie.search('in'))
    # (True, 1)
    print(trie.search('inn'))
三分探索木 三分探索木(さんぶんたんさくぎ、英:ternary search tree)は、トライ木の各ノードを二分探索木として表現したデータ構造である。 出典: フリー百科事典『ウィキペディア(Wikipedia)』 三分探索木は、トライ木...