Python自定义决策树：图像化输出及可读文本打印方法咨询

算力阿峰

2026-6-23

决策树可视化与文本打印方案

问题背景

我已经用Python实现了一个以Node类作为节点的决策树，代码如下：

class Node:
    '''
    Helper class which implements a single tree node.
    '''
    def __init__(self, feature=None, threshold=None, data_left=None, data_right=None, gain=None, value=None):
        self.feature = feature
        self.threshold = threshold
        self.data_left = data_left
        self.data_right = data_right
        self.gain = gain
        self.value = value

我想给这个树添加打印功能，输出每个节点的gain、特征名称、threshold，以及叶节点的最终标签value。现在有两个问题：

有哪些库可以把训练后的决策树输出为图像格式？怎么用？
如果不用图像格式，用什么算法能实现包含上述信息的可读性强的文本打印？

我自己写的print_tree方法可读性很差，代码如下：

def print_tree(self,node,depth=0):
    if node is None:
        return
    prefix = "    " * depth
    # If the node is a leaf node, print its value
    if node.value is not None:
        print(f"{prefix}Value: {node.value}")
    else:
        # Print the feature and threshold for the split at this node
        print(f"{prefix}Feature: {node.feature}, Threshold: {node.threshold}")
        # Recursively print the left and right subtrees
        print(f"{prefix}--> Left:")
        self.print_tree(node.data_left, depth + 1)
        print(f"{prefix}--> Right:")
        self.print_tree(node.data_right, depth + 1)

解决方案

一、生成决策树图像的库及使用方法

1. Graphviz

这是最常用的决策树可视化工具，需要同时安装Python包和系统级Graphviz工具：

安装步骤：
1. 安装Python包：
```
pip install graphviz
```
2. 安装系统工具：
  - Ubuntu/Debian：sudo apt install graphviz
  - macOS：brew install graphviz
  - Windows：下载Graphviz安装包并添加到系统PATH

使用示例：

from graphviz import Digraph

def tree_to_graph(node, graph=None):
    # 初始化图
    if graph is None:
        graph = Digraph()
    # 生成节点标签
    if node.value is not None:
        node_label = f"Leaf Node\nValue: {node.value}"
    else:
        node_label = f"Gain: {node.gain:.4f}\nFeature: {node.feature}\nThreshold: {node.threshold:.4f}"
    # 添加当前节点
    graph.node(name=str(id(node)), label=node_label)
    # 递归添加左子树及边
    if node.data_left is not None:
        graph.edge(str(id(node)), str(id(node.data_left)), label="<=")
        tree_to_graph(node.data_left, graph)
    # 递归添加右子树及边
    if node.data_right is not None:
        graph.edge(str(id(node)), str(id(node.data_right)), label=">")
        tree_to_graph(node.data_right, graph)
    return graph

# 生成并保存图像
root_node = # 你的决策树根节点
tree_graph = tree_to_graph(root_node)
tree_graph.render("decision_tree", format="png")  # 保存为PNG文件
tree_graph.view()  # 打开预览窗口

该方法会生成标准的决策树图像，节点包含所有需要的信息，边标注分支条件。

2. Plotly

适合生成交互式决策树图像，无需额外安装系统工具：

安装：
```
pip install plotly
```

使用示例：

import plotly.graph_objects as go

def collect_tree_data(node, parent_id=None, nodes=[], edges=[]):
    node_id = str(id(node))
    # 添加节点信息
    if node.value is not None:
        nodes.append(dict(
            id=node_id,
            label=f"Leaf Node\nValue: {node.value}"
        ))
    else:
        nodes.append(dict(
            id=node_id,
            label=f"Gain: {node.gain:.4f}\nFeature: {node.feature}\nThreshold: {node.threshold:.4f}"
        ))
    # 添加边信息
    if parent_id is not None:
        edges.append(dict(
            from=parent_id,
            to=node_id,
            label="<=" if node == node.data_left else ">"
        ))
    # 递归处理子节点
    if node.data_left:
        collect_tree_data(node.data_left, node_id, nodes, edges)
    if node.data_right:
        collect_tree_data(node.data_right, node_id, nodes, edges)
    return nodes, edges

# 生成交互式图
root_node = # 你的决策树根节点
nodes, edges = collect_tree_data(root_node)
fig = go.Figure(go.Sankey(
    node=dict(
        pad=15,
        thickness=20,
        label=[n["label"] for n in nodes],
        color="blue"
    ),
    link=dict(
        source=[nodes.index(n) for n in nodes if any(e["from"] == n["id"] for e in edges)],
        target=[nodes.index(n) for n in nodes if any(e["to"] == n["id"] for e in edges)],
        label=[e["label"] for e in edges],
        value=[1]*len(edges)
    )
))
fig.update_layout(title_text="Interactive Decision Tree", font_size=10)
fig.show()

生成的图像支持缩放、hover查看节点详情，适合网页展示或交互式分析。

二、可读性强的文本打印实现

可以通过优化递归打印逻辑，使用树形层级符号（如├、└、│）来区分节点层级，同时完整展示gain、特征、阈值等信息：

改进后的代码：

def print_tree(self, node, depth=0, prefix=""):
    if node is None:
        return
    # 处理叶节点
    if node.value is not None:
        print(f"{prefix}└── 叶节点: Value={node.value}")
        return
    # 处理内部节点，展示所有关键信息
    node_info = f"├── 分裂节点: 特征={node.feature}, 阈值={node.threshold:.4f}, 增益={node.gain:.4f}"
    print(f"{prefix}{node_info}")
    # 左子树前缀：当前层级有右兄弟节点，保留垂直连线符号
    left_prefix = f"{prefix}│   "
    print(f"{left_prefix}左分支 (<= 阈值):")
    self.print_tree(node.data_left, depth + 1, left_prefix + "    ")
    # 右子树前缀：无后续兄弟节点，用空格替代垂直连线
    right_prefix = f"{prefix}    "
    print(f"{right_prefix}右分支 (> 阈值):")
    self.print_tree(node.data_right, depth + 1, right_prefix + "    ")

打印效果示例：

├── 分裂节点: 特征=age, 阈值=30.0000, 增益=0.8921
│   左分支 (<= 阈值):
│       └── 叶节点: Value=0
    右分支 (> 阈值):
        ├── 分裂节点: 特征=income, 阈值=50000.0000, 增益=0.6789
        │   左分支 (<= 阈值):
        │       └── 叶节点: Value=1
            右分支 (> 阈值):
                └── 叶节点: Value=0

这种打印方式通过树形符号清晰展示节点层级关系，每个节点的关键信息完整呈现，比原代码可读性提升明显。

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

Python自定义决策树：图像化输出及可读文本打印方法咨询

问题背景 #

解决方案 #

一、生成决策树图像的库及使用方法 #