第三章 决策树 3.2 使用 Matplotlib 注解绘制树形图

上节学习了如何从数据集中创建树，但是字典的表示形式非常不易于理解，而且直接绘制图形也比较困难，这一节我们将使用 Matplotlib库 来创建树形图。

3.2.1 Matplotlib注解

Matplotlib提供了一个注解工具 annotations，可以在数据图形上添加文本注释。

创建一个文件，命名为 treePlotter.py ，然后输入：

# -*- coding:utf-8 -*-
import matplotlib.pyplot as plt


# 定义文本框和箭头格式
decisionNode = dict(boxstyle = "sawtooth", fc = "0.8") # fc 应该是颜色深浅
leafNode = dict(boxstyle = "round4", fc = "0.8")
arrow_args = dict(arrowstyle = "<-")


def plotNode(nodeTxt, centerPt, parentPt, nodeType):
    # centerPt 箭头指向坐标， parentPt 箭头终点坐标
    createPlot.ax1.annotate(nodeTxt, xy = parentPt,\
    xycoords = 'axes fraction',\
    xytext = centerPt, textcoords = 'axes fraction',\
    va = "center", ha = "center", bbox = nodeType, arrowprops = arrow_args)


def createPlot():
    fig = plt.figure(1, facecolor = 'white')
    fig.clf()
    createPlot.ax1 = plt.subplot(111, frameon = False)
    plotNode(U'决策节点', (0.5, 0.1), (0.1, 0.5), decisionNode) # U 这里指的是 utf 编码
    plotNode(U'叶节点', (0.8, 0.1), (0.3, 0.8), leafNode)
    plt.show()

建立一个运行文件 run_treePlotter.py ，输入：

# run_treePlotter.py
import treePlotter
print '>>> treePlotter.createPlot()'
treePlotter.createPlot()

结果如下：