1.需求说明
求出决策树的字典存储形式数据后,绘制出决策树的图形,则会更形象认识和了解其决策树。
比如,有决策树的字典结构如下所示:
tree_dict = {'house?': {'hourse_no': {'working?': {'work_no': 'refuse', 'work_yes': 'agree'}}, 'hourse_yes': 'agree'}}
要绘制出对应的如下决策树:
本章代码就是完成此需求的。
2. 代码
# -*- coding: utf-8 -*-
"""
@author: 蔚蓝的天空Tom
Aim:得到决策树的字典后,需要使用python来绘制对应的决策树figure
输入决策树的字典,样例如下所示:
dtree = {'house?': {'hourse_no': {'working?': {'work_no': 'refuse', 'work_yes': 'agree'}}, 'hourse_yes': 'agree'}}
"""
import matplotlib.pyplot as plt
#定义判断结点形状,其中boxstyle表示文本框类型,fc指的是注释框颜色的深度
decisionNode = dict(boxstyle="round4", color='r', fc='0.9')
#定义叶结点形状
leafNode = dict(boxstyle="circle", color='m')
#定义父节点指向子节点或叶子的箭头形状
arrow_args = dict(arrowstyle="<-", color='g')
def plot_node(node_txt, center_point, parent_point, node_style):
'''
绘制父子节点,节点间的箭头,并填充箭头中间上的文本
:param node_tx