Python自然语言处理之字典树知识总结

什么是字典树

字典树（Trie树）是一种哈希树的变种，也称为单词查找树或键树。字典树是一棵树，每个节点包含若干字符，而不是单个字符。在实现自然语言处理中，字典树常用来处理字符串匹配、拼写检查、词频统计等任务。

字典树的优势在于它可以在$O(k)$的时间复杂度($k$为字符串长度)内完成字符串的查找操作，而且还可以较方便的实现自动补全功能。

字典树的实现

字典树的实现可以通过Python中的类来进行。一个字典树的节点需要记录三个信息：当前字符、是否形成一个单词、以当前节点为前缀的单词数目。

下面是一个基本字典树类的定义：

class TrieNode:
    def __init__(self):
        self.children = [None] * 26
        self.is_end_of_word = False
        self.word_count = 0

其中，children列出了所有可能的子节点，这里用字符的ASCII码范围来作为节点个数，即26个字母。is_end_of_word表示从根节点到当前节点是否形成一个单词。word_count表示以当前节点为前缀的单词个数。

下面是一个完整字典树的实现：

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

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

    def insert(self, word):
        node = self.root
        for char in word:
            idx = self._char_to_index(char)
            if not node.children[idx]:
                node.children[idx] = TrieNode()
            node = node.children[idx]
            node.word_count += 1
        node.is_end_of_word = True

    def search(self, word):
        node = self.root
        for char in word:
            idx = self._char_to_index(char)
            if not node.children[idx]:
                return False
            node = node.children[idx]
        return node.is_end_of_word

    def starts_with(self, prefix):
        node = self.root
        for char in prefix:
            idx = self._char_to_index(char)
            if not node.children[idx]:
                return False
            node = node.children[idx]
        return True

    def count_starts_with(self, prefix):
        node = self.root
        for char in prefix:
            idx = self._char_to_index(char)
            if not node.children[idx]:
                return 0
            node = node.children[idx]
        return node.word_count

其中，_char_to_index(char)用于将字符映射为下标，与节点的children数组对应。insert(word)用于插入一个单词，search(word)用于查找一个单词是否存在，starts_with(prefix)用于查找所有以给定前缀开头的单词，count_starts_with(prefix)用于统计所有以给定前缀开头的单词数目。

字典树的应用举例

示例一：拼写检查

首先，构建一个字典树并加入常用单词：

trie = Trie()
with open('english_words.txt', encoding='utf-8') as f:
    for word in f:
        trie.insert(word.strip())

然后，对给定文本进行拼写检查：

def spell_check(text):
    words = re.findall(r'\w+', text.lower())
    for word in words:
        if not trie.search(word):
            print(f'"{word}" is misspelled')

这里用正则表达式r'\w+'将文本中的单词提取出来，并将每个单词转化为小写。然后，遍历每个单词，若不在字典树中，则表示该单词拼写错误。

示例二：自动补全

首先，给字典树加入一些单词：

trie = Trie()
words = ['apple', 'banana', 'cherry', 'grape']
for word in words:
    trie.insert(word)

然后，定义一个函数来返回所有以给定前缀prefix开头的单词：

def autocomplete(trie, prefix):
    node = trie.root
    for char in prefix:
        idx = trie._char_to_index(char)
        if not node.children[idx]:
            return []
        node = node.children[idx]

    res = []
    def recurse(node, prefix):
        if node.is_end_of_word:
            res.append(prefix)
        for i, child in enumerate(node.children):
            if child:
                recurse(child, prefix + chr(i + ord('a')))          
    recurse(node, prefix)
    return res

该函数首先找到与前缀prefix匹配的节点，然后使用递归的方式从该节点开始遍历整个字典树，当遇到一个单词结束时，保存该单词。

使用该函数进行自动补全：

words = autocomplete(trie, 'a')
print(words)

结果为['apple']，表示所有以'a'开头的单词中，只有'apple'一种。