使用graphviz绘制决策树更好一点,可以避免节点重叠,并且导出的pdf属于非图片型pdf。
绘制决策树大体分为3个步骤:
1.准备数据
2.创建并训练决策树模型
3.设置决策树参数绘制决策树
示例代码如下:
python
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()