数据可视化matplotlib(03) 绘制决策树

简介

决策树的主要优点是直观易于理解，如果不能将其直观的显示出来，就无法发挥其优势。本文将使用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]