简介
决策树的主要优点是直观易于理解,如果不能将其直观的显示出来,就无法发挥其优势。本文将使用matplotlib来绘制树形图,并讲解具体的代码实现。
决策树图实例
为了便于理解,我们先来看看实际的决策树的图长个什么样子。下图所示的流程图就是一个决策树,正方形代表判断模块(decision block), 椭圆形代表终止模块(terminal block),表示已经得出结论,可以终止运行。
代码实现
这里我们给出了相关的代码实现,后面会对一些重要的实现进行详解。
# /usr/bin/python
# -*- coding: UTF-8 -*-
'''
Created on 2017年11月16日
@author: bob
'''
import matplotlib.pyplot as plt
# pylint: disable=redefined-outer-name
# 定义文本框和箭头格式
decision_node = dict(boxstyle="sawtooth", fc="0.8")
leaf_node = dict(boxstyle="round4", fc="0.8")
arrow_args = dict(arrowstyle="<-")
def retrieve_tree(i):
list_of_trees = [{'no surfacing': {0: 'no', 1: {'flippers': {0: 'no', 1: 'yes'}}}},
{'no surfacing': {0: 'no', 1: {'flippers': {0: {'head': {0: 'no', 1: 'yes'}}, 1: 'no'}}}}
]
return list_of_trees[i]
def get_num_leafs(mytree):
'''
获取叶子节点数
'''
num_leafs = 0
first_str = mytree.keys()[0]
second_dict = mytree[first_str]
for key in second_dict.keys():
if type(second_dict[key]).__name__ == 'dict':
num_leafs += get_num_leafs(second_dict[key])
else:
num_leafs += 1
return num_leafs
def get_tree_depth(mytree):
'''
获取树的深度
'''
max_depth = 0
first_str = mytree.keys()[0]