如何在Julia的DecisionTree.jl中为决策树打印添加特征名称

问题描述 #

使用DecisionTree.jl的print_tree函数打印决策树时，默认输出Feature X格式的特征标识，希望替换为自定义的特征名称（如“花萼长度”“花瓣宽度”等）。

解决方案 #

官方print_tree函数未直接支持传入自定义特征名称，可通过实现自定义打印函数来实现需求。步骤如下：

1. 纠正原代码中的拼写错误（feauture_names改为feature_names）
2. 编写自定义的print_tree_with_names函数，遍历决策树节点并替换特征名称
3. 调用自定义函数完成打印

修改后的完整代码 #

using DecisionTree

# 加载鸢尾花数据集
features, labels = load_data("iris")
features = float.(features)
labels   = string.(labels)

# 定义自定义特征名称
feature_names = ["花萼长度", "花萼宽度", "花瓣长度", "花瓣宽度"]

# 训练深度为2的分类器
model = DecisionTreeClassifier(max_depth=2)
fit!(model, features, labels)

# 自定义打印函数：支持替换特征名称
function print_tree_with_names(tree::Union{DecisionTreeClassifier, Node}, feature_names::Vector{String}, depth::Int=5, indent::Int=0)
    prefix = "    "^indent
    if tree isa Node
        # 替换特征ID为自定义名称
        feature_name = feature_names[tree.featid]
        println("$(prefix)\"$(feature_name)\", Threshold $(tree.featval)")
        # 递归打印左右子节点（控制深度）
        indent < depth && print_tree_with_names(tree.left, feature_names, depth, indent+1)
        indent < depth && print_tree_with_names(tree.right, feature_names, depth, indent+1)
    else
        println("$(prefix)$(tree)")
    end
end

# 调用自定义函数打印决策树
print_tree_with_names(model.root, feature_names, 5)

输出结果 #

"花瓣长度", Threshold 2.45
    L-> Iris-setosa : 50/50
    R-> "花瓣宽度", Threshold 1.75
        L-> Iris-versicolor : 49/54
        R-> Iris-virginica : 45/46

说明 #

自定义函数通过访问模型的root节点（决策树的根节点），遍历每个节点的featid属性，将其映射为feature_names中对应的名称，同时保留了原print_tree的缩进和深度控制逻辑，确保输出格式与原函数一致。

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