上节学习了如何从数据集中创建树,但是字典的表示形式非常不易于理解,而且直接绘制图形也比较困难,这一节我们将使用 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()
结果如下: