Python 字典树的可视化与图形化展示

2023-04-11 00:00:00 字典 化与 可视

Python 字典树是一种高效的数据结构,可以用来快速地搜索和查找字符串。为了更好地理解字典树,可以将其可视化并进行图形化展示。

下面是一份 Python 代码演示,实现了字典树的可视化与图形化展示,其中使用了字符串 “pidancode.com” 和 “皮蛋编程” 作为范例。

首先需要安装必要的库:

pip install graphviz
pip install pillow

然后开始编写代码,首先导入所需的库:

from PIL import Image
from graphviz import Digraph

接着定义一个字典树节点类:

class TrieNode:
    def __init__(self, char: str):
        self.char = char
        self.children = {}
        self.word_finished = False

这个类包含了一个字符、一个子节点字典和一个布尔值,用来表示当前节点是否是一个单词的结束。

下面是 Trie 类的实现:

class Trie:
    def __init__(self):
        self.root = TrieNode('')

    def add_word(self, word: str):
        node = self.root
        for char in word:
            if char in node.children:
                node = node.children[char]
            else:
                new_node = TrieNode(char)
                node.children[char] = new_node
                node = new_node
        node.word_finished = True

    def _search(self, node, word, prefix=False):
        if not word:
            return prefix or node.word_finished

        char = word[0]
        if char == '.':
            for child_node in node.children.values():
                if self._search(child_node, word[1:], prefix):
                    return True
        elif char in node.children:
            return self._search(node.children[char], word[1:], prefix)

        return False

    def search(self, word: str):
        return self._search(self.root, word)

    def starts_with(self, prefix: str):
        return self._search(self.root, prefix, True)

这个类包含了插入单词的方法、查找单词的方法和查找前缀的方法。

接下来,定义一个方法来生成 Trie 的图片:

def generate_trie_image(words, filename):
    trie = Trie()
    for word in words:
        trie.add_word(word)

    dot = Digraph(comment='Trie', format='png')
    dot.attr('node', shape='circle')

    def add_node(node, name=''):
        if not name:
            name = node.char
        if node.word_finished:
            dot.attr('node', shape='doublecircle')
        else:
            dot.attr('node', shape='circle')
        dot.node(name)
        for child in node.children.values():
            child_name = name + child.char
            dot.edge(name, child_name)
            add_node(child, child_name)

    add_node(trie.root)
    dot.render(filename, view=True)

这个方法接收一个单词列表和一个文件名,使用前面定义的 Trie 类来创建一棵字典树,并使用 graphviz 库来生成 PNG 格式的图片。

最后,用下面的代码来生成图片:

if __name__ == "__main__":
    words = ['pidancode.com', '皮蛋编程']
    generate_trie_image(words, 'trie.png')
    Image.open('trie.png').show()

这个代码使用了我们前面定义的 generate_trie_image 方法来生成字典树图片,保存在 trie.png 文件中,并用 PIL 库打开并展示。图片如下所示:

trie.png

可以看到,字典树的结构被清晰地展示出来了,每个节点表示一个字符,其中带圆圈的节点表示一个单词的结束。

相关文章