数据结构 | 30行代码，手把手带你实现Trie树

今天是算法和数据结构专题的第28篇文章，我们一起来聊聊一个经典的字符串处理数据结构——Trie。

在之前的4篇文章当中我们介绍了关于博弈论的一些算法，其中应用最广也是最重要的就是最后的SG函数。了解到这些之后，足够我们应付常见的博弈论算法问题了。博弈论本身就是一门学科，其中有着很深邃的理论基础，我们只是浅尝辄止，大家感兴趣的可以自行钻研一下，相信一定会很有收获。

小故事

以前读过一个大牛的文章，文章里讨论了一个问题，如果不是为了面试的话，我们为什么要学算法？

他讲了一个他自己的故事，说是在很多年前，手机还是诺基亚功能机的时代，他为塞班系统开发了一个通讯簿查找联系人的软件。软件的功能很简单，就是存储联系人，然后可以通过拼音或者是拼音首字母查找到对应的联系人。这里需要对汉字以及拼音的映射做一个处理，也不是很复杂的操作，我们脑补应该就可以想出来。

软件很快做好了，做好了之后投入使用发现也很好用。但是很快遇到了一个没想到的问题，就是当联系人多了之后，软件的运行速度变得非常慢，也就是卡。卡的原因也很简单，因为搜索联系人的这个步骤他用的是遍历查找的方式搜索的。他一开始先是自己脑补了一些优化方案和野路子，虽然能有些提升但是不能根本解决问题。后来被逼无奈，他在搜索了相关资料之后，找到了我们今天的主角Trie，用上了这个算法之后，这个问题瞬间迎刃而解，即使存储了成千上万的联系人再也不会卡顿了。从此他大彻大悟，算法并不是奇淫技巧，真的是有用的。

我们就以这个文章当中的问题作为基础，来看看Trie的原理，以及它为什么可以解决这个问题。

Trie简介

Trie树有好几个中文翻译名字，有的称为字典树，有的称为前缀树。这都是可以的，看大家各位的喜欢。我一般就称Trie树，对方听不懂才会说字典树XD。

从字典树和前缀树的称谓当中我们是可以脑补出来它的大概原理的，也就是以字典和前缀的形式存储数据。听着有点抽象，其实我们看一张图就完全明白了。

从图中我们可以看出来，Trie树是一棵多叉树。树中的每一个节点存储一个字符，我们从根节点到节点的路径上的字符连起来就成了单词。也就是说所有的单词都是这样纵向的形式存储在树上的。

这样存储有什么好处呢？最大的好处就是拥有相同前缀的单词可以共享前缀，比如ana，ann和and这两个单词前两个字符是相同的都是an，所以他们拥有一条公共前缀链路。在这样的结构之下，当我们需要查找单词是否在树中的时候，我们只需要从树根开始遍历，如果能找到对应的结尾节点就说明单词存在，否则就说明单词不存在。

举个例子，比如我们要查找单词doe，我们先从根节点查找字符d。发现d存在，于是转移到了d。接着我们查找字符o，在d节点下面也找到了字符o，转移到o，查找字符e，发现e不存在，于是就说明了doe单词不存在。但是这里有一个问题，假设我们存在一个单词是doea，我们查找doe还是可以找到，但是doe单词其实是不存在的，这不就错误了吗？

的确，这样可能存在问题，所以我们需要在节点当中记录一下，是否是某一个单词的结尾。这样我们不仅需要找到对应的单词，还需要防止我们将其他单词的前缀当成是单词。

我们插入单词的过程和查询非常接近，同样是一个树上遍历的过程，只不过如果我们发现查询的节点不存在时会手动创建。整个单词插入完成之后，将最后一个字符对应的节点进行标记，表明这是一个单词的结尾。

简单的Trie树只需要完成添加和查询即可，如果要涉及删除，我们只需要在节点当中维护一下经过该节点的单词数量即可。在删除的时候，将沿途经过的节点标记的数量-1，如果遇到数量为0的节点，直接删除即可。

代码实现

光说不练假把式，我们自然也是要来练练的。

相信大家也从描述当中看出来了，Trie的原理说穿了其实很简单，实现起来也不困难。网上有许多版本，很多是面向过程实现的，我把它封装了一下，用Python面向对象实现了一个版本。理解了原理之后，大家可以根据自己的需要开发自己的风格的版本，代码其实不太重要，主要还是理解原理。

我把Trie树分成了两个部分，第一个部分是树上的节点。对于Trie树上的节点来说它需要提供两个功能。第一个功能是返回当前节点是否是某一个字符串的结尾，第二个功能是根据字符查找后继的节点。我们只需要在类当中设置一个flag标记和一个dict属性来存储后继元素就行了。

class Node:
    def __init__(self, is_leaf=False):
        self.child = {}
        self._is_leaf = is_leaf

    @property
    def is_leaf(self):
        return self._is_leaf

    @is_leaf.setter
    def is_leaf(self, is_leaf):
        self._is_leaf = is_leaf

    # 加入孩子节点
    def put(self, key, value):
        self.child[key] = value

    @staticmethod
    # 将字符转化成数字
    # 其实没有必要，因为用到了dict，如果用数组存储孩子的话，需要用它来计算下标
    def get_idx_of_str(_str):
        if len(_str):
            return -1
        if ord('a') <= ord(_str) <= ord('z'):
            return ord(_str) - ord('a')
        else:
            return ord(_str) - ord('A') + 26

    # 根据字符获取下一个节点
    def get_next_node(self, _str):
        if len(_str) != 1:
            return None
        idx = Node.get_idx_of_str(_str)
        return self.child.get(idx, None)

    def get_node(self, key):
        return self.child.get(key, None)

这里我将is_leaf两个方法用property封装，从而可以方便使用，这个也是常用的惯例。有了节点之后，我们再开发Trie类就很方便了，对于Trie这个类而言我们只需要实现两个方法，一个是插入字符串，一个是字符串的查询。在有了Node类之后，这两个方法实现也很简单了。

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

    def insert(self, _str):
        cur = self.root
        # 遍历字符
        for c in _str:
            # 查找下一个节点
            if cur.get_next_node(c) is None:
                # 如果节点不存在，自己创建一个新节点并插入
                key = Node.get_idx_of_str(c)
                cur.put(key, Node())
                cur = cur.get_node(key)
            # 否则继续往下
            else:
                cur = cur.get_next_node(c)
        cur.is_leaf = True
    
    def query(self, _str):
        cur = self.root
        # 遍历字符
        for c in _str:
            # 查询，如果查询不到返回False
            if cur.get_next_node(c) is None:
                return False
            cur = cur.get_next_node(c)
        # 返回是否是字符串结尾
        return cur.is_leaf

这两段代码应该都不能读懂，最后，我们尝试一下使用它来测试一下：

if __name__ == "__main__":
    trie = Trie()
    trie.insert('abcda')
    trie.insert('abcde')
    trie.insert('eecdab')
    trie.insert('mout')
    trie.insert('ymm')
    print(trie.query('abcda'))
    print(trie.query('mout'))
    print(trie.query('ym'))

输出的结果和我们预期一致，说明大概率是正确的。

总结

Trie树中我们将字符串相同的前缀存储在了同样的链路上，节省了大量空间的消耗。并且在查询单词的时候，我们沿着Trie树进行遍历，只需要单词长度的时间就可以得到结果。并且我们可以在Node这个类当中存储其他一些我们需要的信息，这样Trie就转化成了一个以string为key的dict。

Trie树在机器学习领域当中应用也非常广泛，尤其是自然语言处理。可以实现文本的快速分词、词频统计、模糊匹配等功能。并且Trie树还有很多拓展，比如压缩数据空间的双数组Trie树以及AC自动机等等。