使用graphviz绘制决策树更好一点,可以避免节点重叠,并且导出的pdf属于非图片型pdf。

绘制决策树大体分为3个步骤:

1.准备数据
2.创建并训练决策树模型
3.设置决策树参数绘制决策树

示例代码如下:

from sklearn import tree
import graphviz

#1.准备数据
x=[[2, 0, 0, 2],
   [2, 1, 1, 0],
   [2, 1, 1, 2],
   [2, 1, 1, 1],
   [2, 0, 1, 0],
   [2, 0, 1, 1],
   [2, 0, 0, 1],
   [0, 1, 1, 0],
   [0, 1, 1, 1],
   [0, 1, 0, 0],
   [0, 1, 0, 1],
   [0, 0, 1, 2],
   [0, 0, 1, 1],
   [1, 1, 0, 0],
   [1, 1, 0, 2],
   [1, 0, 1, 0],
   [1, 0, 1, 2],
   [1, 0, 1, 1],
   [1, 0, 0, 0],
   [1, 0, 0, 1]]
y = [1, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0]

#设置类别映射表
class_map = {
    0: '不满意',
    1: '满意'
}

#2.创建并训练决策树模型
model = tree.DecisionTreeClassifier()
model.fit(x, y)

#3.设置决策树参数绘制决策树
dot_data = tree.export_graphviz(
    model,
    feature_names=['维护类型','环境温度','环境湿度','通风效能'],
    class_names = list(map(lambda x:class_map[x],model.classes_)),
    filled = True,
    rounded = True,
    fontname='SimHei',
    special_characters=True,
    leaves_parallel=False

)
graph = graphviz.Source(dot_data)
graph.render("DecisionTree", format="png")
graph.view()

更多推荐