如何改善我的Trie实现的初始化？

Question

我正在尝试从一个巨大的单词列表中读取单词，并以一种可以快速检索的方式存储它们。起初我想使用一种叫做“字典树”的结构，老实说，我的实现方式有点简单，就是用嵌套的哈希表，每个键代表一个不同的字母。目前，把一个单词插入到字典树里需要很长时间（运行这个程序需要20秒以上），我想知道有没有人有什么建议可以帮助我提高插入的速度？这不是作业。

import string
import time

class Trie:

    def __init__(self):
        self.root = TrieNode()

    def insert_word(self, word):
        current_node = self.root
        for letter in word:
            trie_node = current_node.get_node(letter)
            current_node = trie_node

class TrieNode:

    def __init__(self):
        self.data = {}

    def get_node(self, letter):
        if letter in self.data:
            return self.data[letter]
        else:
            new_trie_node = TrieNode()
            self.data[letter] = new_trie_node
            return new_trie_node

def main():
    start_time = time.time()
    trie = Trie()

    with open('/usr/share/dict/words', 'r') as dictionary:
        word_list = dictionary.read()
    word_list = word_list.split("\n")

    for word in word_list:
        trie.insert_word(word.lower())

    print time.time() - start_time, "seconds"


if __name__ == "__main__":
    main()

Answer 1

这里有一个Trie的替代实现，可以查看。

比较一下 Trie.insert_word 和 build：

def build(words,trie={'end':False}):
    '''
    build builds a trie in O(M*L) time, where
        M = len(words)
        L = max(map(len,words))
    '''
    for word in words:
        pt=trie
        for letter in word:
            pt=pt.setdefault(letter, {'end':False})
        pt['end']=True
    return trie

在 Trie 中，对于每个 word 中的 letter，insert_word 会调用 current_node.get_node(letter)。这个方法里有一个 if 和一个 else 的分支，每当走到 else 的时候，就必须创建一个新的 TrieNode，然后把新的键值对插入到 self.data 这个字典里。

而在 build 中，trie 本身就是一个字典。对于 word 中的每个 letter，只需要调用一次 pt.setdefault(...)。字典的方法是用C语言实现的，比用Python写类似的代码要快。

timeit 显示出大约有2倍的速度差异（build 更快）：

def alt_main():
    with open('/usr/share/dict/words', 'r') as dictionary:
        word_list = dictionary.read()
    word_list = word_list.split("\n")
    return build(word_list)


% python -mtimeit -s'import test' 'test.main()'
10 loops, best of 3: 1.16 sec per loop

% python -mtimeit -s'import test' 'test.alt_main()'
10 loops, best of 3: 571 msec per loop

Answer 2

在你确认搜索功能是否正常之前，努力加快你的字典树初始化是毫无意义的。

在@unutbu提到的代码中，你觉得为什么要搞那些 {'end':False} 和 pt['end']=True 的操作呢？

这里有一些测试数据给你参考：

words_to_insert = ['foo', 'foobar']
queries_expecting_true = words_to_insert
queries_expecting_false = "fo foe foob food foobare".split()

还有一个想法：你没有说明你想要的不仅仅是判断一个查询词是否存在。如果是这样的话，你应该考虑把你自己做的字典树和内置的 set 进行对比。比较的标准包括：加载速度（可以考虑从一个存档文件中加载），查询速度，以及内存使用情况。

如果你想要获取的信息不仅仅是一个 bool 值，那就把 dict 替换成 set，然后再看看这个回答。

如果你想在输入字符串中搜索单词，那么你可以考虑@unutbu提到的代码，修复一些错误，并在 find 函数中进行一些加速（只计算一次 len(input)，使用 xrange 代替 range（Python 2.x）），并去掉不必要的 TERMINAL: False 条目：

TERMINAL = None # Marks the end of a word

def build(words, trie=None): # bugs fixed
    if trie is None:
        trie = {}
    for word in words:
        if not word: continue # bug fixed
        pt = trie # bug fixed
        for ch in word:
            pt = pt.setdefault(ch, {})
        pt[TERMINAL] = True
    return trie

def find(input, trie):
    len_input = len(input)
    results = []
    for i in xrange(len_input):
        pt = trie
        for j in xrange(i, len_input + 1):
            if TERMINAL in pt:
                results.append(input[i:j])
            if j >= len_input or input[j] not in pt:
                break
            pt = pt[input[j]]
    return results    

或者你可以看看 Danny Yoo 的快速实现，这是 Aho-Corasick 算法 的实现。