Trie树-提升海量数据的模糊查询性能

今天这篇文章源于上周在工做中解决的一个实际问题，它是个比较广泛的问题，不管作什么开发，估计都有遇到过。具体是这样的，咱们有一份高校的名单（2657个），须要从海量的文章标题中找到包含这些高校的标题，其实就是模糊查询。（关注公众号 渡码，回复关键词 trie 获取完整源代码）对应的伪代码以下

selected_titles = []
for 标题 in 海量标题：
  for 高校 in 高校名单：
    if 标题.contains(高校):
      selected_titles.add(标题)
      break复制代码

若是是大数据开发，对应的SQL的伪代码是这样的

select title
from tb
where 
  title rlike '清华大学|北京大学|...2657个高校'复制代码

上面这两种作法都能实现咱们的需求，但它们的共同问题是查询效率过低。若是咱们要匹配的高校不是2657个而是几十万甚至上百万，那这种方式耗费时之久是不可想象的。

优化这类问题一般须要在数据结构上作文章，这个问题中咱们能优化的数据结构也只有“高校名单”这个了，上面的伪代码中咱们存放“高校名单”的数据结构是数组，当咱们查找某个title是否包含某个高校的时候，须要从头至尾遍历一遍“高校名单”，而且名单越长，遍历耗时就越长。

清楚了数组这种数据结构的缺点后，接下来咱们重点要作的就是寻找一个数据结构能够作到在不遍历整个“高校名单”的状况下就能够完成模糊查询。这个数据结构就是咱们今天要介绍的 Trie 树，冷眼一看这个单词有点陌生，又是一个树型结构，感受会很复杂似的，实际上这个数据结构的设计思想很是简单，一学就会。

下面咱们就来学习一下 Trie 树。为了方面讲解，假设“高校名单”里只有下面5个元素

ABC、ABD、BCD、BCE、C、CAB、CDE复制代码

对应的两种数据结构以下：

抛开这两种数据结构查找的时间复杂度，咱们先从直观上看看为何 Trie 树的查找效率要比数组高。假设咱们要查找，“CDE”这个字符串，在数组结构中，咱们要遍历一遍数组，比较7次才能找到结果，作了比较多的“无用功”。而在 Tire 树中只须要比较3次就能够找到，它的优点很是明显，因为树型结构咱们根本不用考虑左侧A、B开头的两个分支，这就大大减小了比较的次数，从而减小“无用功”。下面用一个动画来演示一下如何建立 Trie 树，以及在 Trie树上查找字符串（若是视频播放不了能够看源码目录中的gif）

树的创建过程其实就是遍历字符串每一个元素并在树上创建相应的节点。字符串查找过程其实就是按照字符串对树进行遍历。Trie 树的创建与字符串查找仍是比较简单的。

不知道你们是否注意到上图 Trie 树中的节点有两种颜色——白色和绿色。绿色节点表明从根节点到当前节点的字符串是“高校名单”中的字符串，也就是咱们创建 Trie 树用到的字符串。以最左侧的叶子结点“C”为例，它表明“ABC”字符串是“高校名单”中的字符串。同理，字符串“AB”就不是“高校名单”里的元素，由于“B”节点不是绿色的，所以当咱们在这棵树上查找字符串“AB”时，是查不到的。这一点须要你们注意，下面编码中咱们也有体现。

另外，有朋友可能会有疑问，咱们最开始的需求不是模糊查询吗，在 Trie 树讲解这部分怎么都在说字符串全词（精确）匹配。这是由于全词匹配是 Tire 树支持的最基本的查找方式，在此基础上，咱们作一些变通就能够很容易实现模糊匹配。

接下来，咱们就来看看代码实现（Python版），首先建立两个数组

colleges = utils.read_file_to_list('key_words.txt')
titles = utils.read_file_to_list('titles.txt')复制代码

colleges就是咱们一直在说的“高校名单”，titles即是“海量标题”，它们都是一维数组，数组每一个元素都是一个字符串。

再来编写 Trie 树相关的代码，若是理解了 Trie 树的设计思想，再编写下面的代码其实很容易。首先要定义一个类表明 Trie 树节点

class TrieNode:
    def __init__(self):
        self.nodes = dict()
        # is_end=True 表明从根节点到当前节点构造Trie树的字符串（出如今“高校名单”里）
        self.is_end = False复制代码

is_end=True就是咱们上面说的绿色节点。

再来编写建立 Trie 树的代码，代码在 TrieNode 类中

def insert_many(self, items: [str]):
    """ 支持输入字符串数组，直接构造一个 Trie 树 :param items: 字符串数组 :return: None """
    for word in items:
        self.insert(word)
​
def insert(self, item: str):
    """ 向 Trie 树插入一个短语 :param item: 待插入的字符串 :return: None """
    curr = self
    for word in item:
        if word not in curr.nodes:
            curr.nodes[word] = TrieNode()
        curr = curr.nodes[word]
    curr.is_end = True复制代码

再来编写查找 Trie 树的代码，代码在 TrieNode 类中

def suffix(self, item: str) -> bool:
    """ 匹配前缀，也就是判断item字符串是不是以“高校名单”中某个字符串开头 :param item: 待匹配字符串 :return: True or False """
    curr = self
    for word in item:
        if word not in curr.nodes:
            return False
        curr = curr.nodes[word]  # 取得子节点
        if curr.is_end:  # 若是is_end=True说明当前字符串包含了“高校名单”的某个字符串
            return True
    return False  # 未匹配上复制代码

这里并非全词匹配，而是前缀匹配，也就是判断待查找的字符串item是不是以“高校名单”中某个字符串开头。

再来编写模糊匹配，代码在 TrieNode 中

def infix(self, item: str) -> bool:
    for i in range(len(item)):
        sub_item = item[i:]  # 将待查找的字符串分红不一样子串
        # 若是子串的前缀在 Trie 树中能匹配上
        # 说明待查找的字符串item中包含“高校名单”中的元素，
        # 即实现了 tile rlike '清华大学|北京大学|...其余大学' 的功能
        if self.suffix(sub_item):
            return True
    return False复制代码

这里其实就是把待查找字符串item分红不一样子串去作前缀匹配，若是子串匹配上，那就说整个字符串item就包含了“高校名单”里面的某个字符串。

最后，咱们运行一下上面的代码，并记录查找时间，与最开始数组结构那一版作个对比。代码以下

# 数组版本
cnt = 0
start_time = int(time.time() * 1000)
for title in titles:
    for x in colleges:
        if x in title:
            cnt += 1
            break
​
end_time = int(time.time() * 1000)
print(cnt)
print('spend: %.2fs' % ((end_time - start_time) / 60.0))
​
# Trie 树版本
root = TrieNode()
root.insert_many(colleges)
cnt = 0
start_time = int(time.time() * 1000)
for title in titles:
    if root.infix(title):
        cnt += 1
​
end_time = int(time.time() * 1000)
print(cnt)
print('spend: %.2fs' % ((end_time - start_time) / 60.0))复制代码

输出结果以下：

5314
spend: 9.13s
5314
spend: 0.23s复制代码

能够看到，用数组匹配用了9s，而用 Trie 树匹配仅用0.23s！

今天介绍的这种提升海量数据模糊查询性能的方式是经过写代码的方式实现的，对于常常写 SQL 的大数据开发者来讲，要把它用起来只是建个 UDF 就能够了，须要在 UDF 的初始化代码中用“高校名单”创建一颗 Trie 树。

今天的内容就分享到这里了，但愿对你有帮助。公众号回复关键词 trie 获取完整源代码

欢迎公众号「渡码」，输出别地儿看不到的干货。