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

本文始发于个人公众号:TechFlow,原创不易,求个关注


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

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

小故事

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

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

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

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

Trie简介

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

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

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

从图中我们可以看出来,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'))
数据结构 | 30行代码,手把手带你实现Trie树

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

总结

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

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

今天的文章到这里就结束了,如果喜欢本文的话,请来一波素质三连,给我一点支持吧(关注、转发、点赞)。

本文使用 mdnice 排版

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

上一篇:小白学PyTorch 动态图与静态图的浅显理解


下一篇:机器学习笔记(五):决策树理论与简单使用