如何使用python构建一棵Trie树，并基于graphviz绘制Trie树

什么是Trie树

Trie树，又叫做字典树，顾名思义它是一种树形的数据结构。Trie树可以做到在一个字符串集合中快速匹配字符串，在构建之后也可以用于词频统计等。

背景

有这样一份数据，它包含了一系列Mac OS系统中的文件路径（实际上记录的是一个进程对于某个文件的操作）

/Users/xxxx/Documents/lll.mp4
/Applications/WeChat.app/Contents/Resources/Compose_Video.svg
/Applications/WeChat.app/Contents/Resources/Compose_Collapsed.svg
/private/var/folders/xxxx

很容易想到的是文件的路径都是由‘/’分割，因此可以把一个文件路径看作Trie树上的一个分支，而文件路径的每一部分都是Trie树中的一个节点，这样我们可以对这份文件路径数据进行路径频次统计，也可以通过绘制整棵Trie树的方式直观的看到当前系统中文件的操作都集中在哪块。

基本的Trie树

虽然场景里并没有用到Trie树的搜索，但是毕竟Trie树的本职还是字符串匹配，这里还是给出完整的Trie树构建和字符串匹配方法。

from collections import defaultdict
from graphviz import Digraph
class TrieNode:
	def __init__(self, data, row):
		self.data = data
		self.children = defaultdict(TrieNode)
		self.count = 1
		self.ending = False  # 节点ending为True则表示当前节点存在数据
		self.id = ''  # 用于标识graphviz图中节点

	def plus(self):
		"""
		用于词频统计
		对于已经存在的数据对节点的count++
		"""
		self.count += 1

class Trie:
	def __init__(self):
		self.root = TrieNode('/')  # 跟节点存储无意义的字符

	def insert(self, items):
		"""
		items为要插入trie树的数据，可以是列表或者字符串
		"""
		p = self.root
		for item in items:
			if item not in p.children:  # item首次出现，创建新节点
				newNode = TrieNode(item)
				p.children[item] = newNode
			else:
				p.children[item].plus()  # 修改当前路径的出现次数
			p = p.children[item]  # 指针指向下一个节点
		p.ending = True

	def find(self, items):
		"""
		在Trie树中查找一个字符串/路径
		"""
		p = self.root
		for item in items:
			if item not in p.children:
				return False
			p = p.children[item]
		if p.ending:
			return True
		else:
			return False  # 节点ending为False则表示此处无数据

paths = ['/Users/xxxx/Documents/lll.mp4', '/Applications/WeChat.app/Contents/Resources/Compose_Video.svg',
		 '/Applications/WeChat.app/Contents/Resources/Compose_Collapsed.svg', '/private/var/folders/xxxx']
trie_data = [path[1:].split('/') for path in paths]
tree = Trie()
for path in trie_data:
	tree.insert(path)

print(tree.find(['Users', 'xxxx', 'Documents', 'lll.mp4']))  # True
print(tree.find(['Users', 'xxxx', 'Documents']))  # False 注意这里可能会对/Users/xxxx/Documents是否存在有歧义，这里认为中间的路径算作不存在

graphviz绘制Trie树

这里省略上面的基本Trie树，仅实现绘制部分的函数。首先使用graphviz绘制Trie树，我们需要绘制所有节点，并连接全部父子节点，这里采用层序遍历的方式来遍历整棵树。

	def level_order(self):
		level_order_res = []  # 层序遍历结果，用于后面连接graphviz图中的节点
		root = self.root
		if root is None:
			return

		g = Digraph('文件路径trie树', node_attr={'color': 'lightblue2'})  # 初始化一张图
		queue = [root]  # 根节点进入队列
		level = 'A'  # 第几层
		idx = 0  # 某层中节点的索引，level与idx一起组成trie树节点的id
		while len(queue) != 0:
			length = len(queue)
			list = []
			for i in range(length):
				id = '{}{}'.format(level, idx)  # 拼接节点id
				node = queue.pop(0)
				node.id = id
				list.append(node)  # 放入当前层的列表中
				g.node(name=id, label=node.data)  # 绘制一个节点

				for key, item in node.children.items():
					queue.append(item)  # 将node的孩子节点都进入队列
				idx += 1  # idx标记向后移动一位
			level_order_res.append(list)
			level = chr(ord(level) + 1)  # level标记进入下一层

		#  绘制边
		for level_nodes in level_order_res:
			for node in level_nodes:
				for key, child_node in node.children.items():
					g.edge(node.id, child_node.id, str(child_node.count))
		g.view()

# 省略上面的Trie树的构建tree
tree.level_order()

总结

其实一个节点不仅可以展示一个路径，也包含了很多其他的数据，我们可以在TrieNode中记录这些数据，并在绘制Trie树时更加多样化，如将更多信息加入节点的label中，也可以通过节点的颜色，节点的形状等来区分。