テーブルによるトライ木の実装

TrieTree.py

class TrieTree:
    def __init__(self):
        size = 256
        self.table = [[None for c in range(size)] for n in range(size)]
        self.node_value = [None for n in range(size)]
        self.last_node = 0

    def insert(self, key, value):
        node = 0
        for c in key:
            next_node = self.table[node][ord(c)]
            if next_node:
                node = next_node
            else:
                self.last_node += 1
                self.table[node][ord(c)] = self.last_node
                node = self.last_node

        self.node_value[node] = value

    def search(self, key):
        node = self._search_node(key)
        return self.node_value[node] if node else None

    def common_prefix_search(self, key):
        node = 0
        d = {}
        for i in range(len(key)):
            next_node = self.table[node][ord(key[i])]
            if next_node:
                if self.node_value[next_node]:
                    d[key[:i + 1]] = self.node_value[next_node]
                node = next_node
            else:
                break

        return d

    def predictive_search(self, key):
        node = self._search_node(key)
        if not node:
            return None
        nodes = [(key, node)]

        d = {}
        while nodes:
            key, node = nodes.pop()
            for c, n in self._get_children(node):
                nodes.append((key + c, n))
                if self.node_value[n]:
                    d[key + c] = self.node_value[n]

        return d

    def _search_node(self, key):
        """
        keyに対応するnode番号を返す
        """
        node = 0
        for c in key:
            next_node = self.table[node][ord(c)]
            if next_node:
                node = next_node
            else:
                return None
        return node

    def _get_children(self, node):
        """
        nodeの(edge, 子供のnode番号)のリストを返す
        """
        return [(chr(i), self.table[node][i]) for i in range(len(self.table[node])) if self.table[node][i]]


def main():
    trie = TrieTree()

    word_value = {"A": 15, "to": 7, "tea": 3, "ted": 4, "ten": 12, "i": 11, "in": 5, "inn": 9}
    for word, value in word_value.items():
        trie.insert(word, value)
    print("検索")
    print({word: trie.search(word) for word in word_value.keys()})
    print({word: trie.search(word) for word in ["t", "te"]})

    print("\n共通接頭辞検索")
    print("inn:", trie.common_prefix_search("inn"))
    print("a:", trie.common_prefix_search("a"))

    print("\n予測検索")
    print("te:", trie.predictive_search("te"))
    print("a:", trie.predictive_search("tea"))
    print("a:", trie.predictive_search("a"))

if __name__ == '__main__':
    main()