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

本文始发于我的公众号：TechFlow，原创不易，求个关注

今天是算法和数据结构专题的第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自动机等等。

今天的文章到这里就结束了，若是喜欢本文的话，请来一波素质三连，给我一点支持吧（关注、转发、点赞）。

本文使用 mdnice 排版