Python3中绘制可视化决策树的可用模块及实现方法咨询

嘿，我来给你梳理几个在Python 3里把这种字典结构转成可视化树形图的好用方法，都是能直接生成图片的，直接上干货：

这俩工具组合是生成树形图的黄金搭档，能导出PNG、PDF等多种格式，生成的图结构规范、可读性强。

首先得安装依赖：

pip install graphviz pydotplus

然后用递归函数遍历你的字典结构，构建可视化图形：

from graphviz import Digraph

def build_tree_graph(node, graph, parent_node=None):
    # 递归遍历字典，添加节点和连线
    for node_name, child_nodes in node.items():
        # 创建当前节点
        graph.node(node_name)
        if parent_node:
            # 连接当前节点和父节点
            graph.edge(parent_node, node_name)
        # 如果子节点是字典，继续递归处理
        if isinstance(child_nodes, dict):
            build_tree_graph(child_nodes, graph, node_name)

# 注意：你原字典里的"comupute"应该是笔误，我改成"compute"啦
MyNodeStructure = {
    'main': {
        'load': {'regA': None},
        'compute': {'loop': {'regA': None, 'regB': None}},
        'store': None
    }
}

# 初始化有向图，指定输出格式为PNG
tree_graph = Digraph(format='png')
# 递归构建图形
build_tree_graph(MyNodeStructure, tree_graph)
# 保存图片到本地（会生成名为decision_tree.png的文件）
tree_graph.render('decision_tree')
# 直接打开预览图形
tree_graph.view()

如果不想依赖Graphviz这种外部工具，treelib是个轻量级的树形结构库，配合PIL可以把文本格式的树转成图片。

先安装依赖：

pip install treelib pillow

然后编写代码：

from treelib import Tree
from PIL import Image, ImageDraw, ImageFont

def add_tree_nodes(tree, parent_id, node_dict):
    for node_name, child_nodes in node_dict.items():
        # 添加当前节点到树中
        tree.create_node(node_name, node_name, parent=parent_id)
        # 递归处理子节点
        if isinstance(child_nodes, dict):
            add_tree_nodes(tree, node_name, child_nodes)

# 初始化树并添加根节点
tree = Tree()
tree.create_node('main', 'main')
# 填充树的结构
add_tree_nodes(tree, 'main', MyNodeStructure['main'])

# 获取树形文本结构
tree_text = tree.show(stdout=False)
# 用PIL把文本转成图片
try:
    # 替换成你系统里的字体路径，比如Windows用'C:/Windows/Fonts/arial.ttf'
    font = ImageFont.truetype('arial.ttf', 14)
except:
    # 找不到字体就用默认的
    font = ImageFont.load_default()

# 计算图片尺寸
lines = tree_text.split('\n')
max_line_width = max(font.getsize(line)[0] for line in lines)
img_width = max_line_width + 20
img_height = len(lines) * 20

# 生成图片
img = Image.new('RGB', (img_width, img_height), color='white')
draw = ImageDraw.Draw(img)
y_offset = 10
for line in lines:
    draw.text((10, y_offset), line, font=font, fill='black')
    y_offset += 20

# 保存图片
img.save('simple_tree.png')

如果需要完全自定义树形图的样式（比如节点颜色、连线样式、布局），可以用Matplotlib手动绘制，灵活性拉满，但需要写更多代码。

安装依赖：

pip install matplotlib

简化版示例代码：

import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle

def plot_custom_tree(node, x, y, node_width, node_height, ax):
    if isinstance(node, dict):
        child_count = len(node.items())
        # 遍历每个子节点
        for idx, (node_name, child) in enumerate(node.items()):
            # 绘制当前节点
            ax.add_patch(Rectangle((x, y - node_height/2), node_width, node_height, fill=False, edgecolor='black'))
            ax.text(x + node_width/2, y, node_name, ha='center', va='center', fontsize=12)
            # 处理子节点
            if child is not None and isinstance(child, dict):
                # 计算子节点的位置
                child_x = x + node_width + 60
                child_y = y - idx * (node_height * 1.8)
                # 绘制连线
                ax.plot([x + node_width, child_x], [y, child_y], color='black', linewidth=1)
                # 递归绘制子树
                plot_custom_tree(child, child_x, child_y, node_width, node_height, ax)
            elif child is None:
                # 叶子节点（值为None）
                child_x = x + node_width + 60
                child_y = y - idx * (node_height * 1.8)
                ax.plot([x + node_width, child_x], [y, child_y], color='black', linewidth=1)
                ax.add_patch(Rectangle((child_x, child_y - node_height/2), node_width, node_height, fill=False, edgecolor='black'))
                ax.text(child_x + node_width/2, child_y, node_name, ha='center', va='center', fontsize=12)

# 创建画布
fig, ax = plt.subplots(figsize=(12, 9))
# 绘制树形图
plot_custom_tree(MyNodeStructure, 50, 400, 80, 40, ax)
# 隐藏坐标轴
ax.axis('off')
# 调整画布范围
ax.set_xlim(0, 700)
ax.set_ylim(0, 450)
# 保存图片
plt.savefig('custom_tree.png', bbox_inches='tight')
plt.show()

内容的提问来源于stack exchange，提问作者Tomato