python 使用文本注解绘制树节点

注解绘制树节点"/>

python 使用文本注解绘制树节点

机器学习实战之决策树(二)在 python 中使用 Matplotlib 注解绘制树形图​blog.csdn

微信公众号：

qiongjian0427qiongjian/Machine-learning​github

运行环境：anaconda—jupyter notebook

Python版本：Python3

1. Matplotlib 注解

Matplotlib 提供了一个注解工具 annotations ，可以在数据图形上添加文本注释。注解通常用于解释数据的内容。

1.1 使用文本注解绘制决策树

import matplotlib.pyplot as plt

decisionNode = dict(boxstyle="sawtooth", fc="0.8") # boxstyle为文本框的类型，sawtooth是锯齿形，fc是边框线粗细

leafNode = dict(boxstyle="round4", fc="0.8") # round4表示圆形

arrow_args = dict(arrowstyle="

def plotNode(nodeTxt, centerPt, parentPt, nodeType):

# nodeTxt:节点文本， centerPt:子节点， parentPt:父节点， nodeType:节点样式

createPlot.axl.annotate(nodeTxt, xy=parentPt,xycoords='axes fraction',\

xytext=centerPt, textcoords='axes fraction',\

va="center", ha="center", bbox=nodeType, arrowprops=arrow_args)

第一个版本的 createPlot()函数。

def createPlot():

fig = plt.figure(1, facecolor='white')

fig.clf()

createPlot.axl = plt.subplot(111, frameon=False)

plotNode('JCJD', (0.5, 0.1), (0.1, 0.5), decisionNode)

plotNode('YJD', (0.8, 0.1), (0.3, 0.8), leafNode)

plt.show()

createPlot 函数是这段代码的核心。

createPlot 函数首先创建了一个新图形并清空绘图区，然后再绘图区上绘制了两个代表不同类型的树节点，后面我们将用这两个节点绘制树图形。

代码运行结果：

2. 构造注解树

定义 getNumLeafs 函数，获取叶节点的数目，getTreeDepth 函数，获取树的层数。

2.1 获取叶节点的数目

def getNumLeafs(myTree):

numLeafs = 0

firstStr = list(myTree.keys())[0]

secondDict = myTree[firstStr]

for key in secondDict.keys():

if type(secondDict[key]).__name__ == 'dict':

numLeafs += getNumLeafs(secondDict[key])

else:

numLeafs += 1

return numLeafs

2.2 获取树的层数

def getTreeDepth(myTree):

maxDepth = 0

firstStr = list(myTree.keys())[0]

secondDict = myTree[firstStr]

for key in secondDict.keys():

if type(secondDict[key]).__name__ == 'dict':

thisDepth = 1 +getTreeDepth(secondDict[key])

else:

thisDepth = 1

if thisDepth > maxDepth : maxDepth = thisDepth

return maxDepth

为了节省时间，函数 retrieveTree 输出预先存储的树信息，避免每次测试都要创建树的麻烦。

def retrieveTree(i):

listOfTrees = [{'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 listOfTrees[i]

运行代码结果：

3. plotTree函数

使用plotMidText函数，计算父节点和子节点的中间位置，并在此处添加简单的文本标签信息。

decisionNode = dict(boxstyle="sawtooth", fc="0.8") # boxstyle为文本框的类型，sawtooth是锯齿形，fc是边框线粗细

leafNode = dict(boxstyle="round4", fc="0.8") # round4表示圆形

arrow_args = dict(arrowstyle="

def plotNode(nodeTxt, centerPt, parentPt, nodeType):

createPlot.ax1.annotate(nodeTxt, xy=parentPt, xycoords='axes fraction',

xytext=centerPt, textcoords='axes fraction',

va="center", ha="center", bbox=nodeType, arrowprops=arrow_args)

def plotMidText(cntrPt, parentPt, txtString):

xMid = (parentPt[0]-cntrPt[0])/2.0 + cntrPt[0] #计算标注位置

yMid = (parentPt[1]-cntrPt[1])/2.0 + cntrPt[1]

createPlot.ax1.text(xMid, yMid, txtString, va="center", ha="center", rotation=30)

计算宽与高

def plotTree(myTree, parentPt, nodeTxt):

numLeafs = getNumLeafs(myTree)

depth = getTreeDepth(myTree)

firstStr = list(myTree.keys())[0] # 找到第一个元素，根节点

cntrPt = (plotTree.xOff + (1.0 + float(numLeafs))/2.0/plotTree.totalW, plotTree.yOff) # 节点位置

plotMidText(cntrPt, parentPt, nodeTxt) #标记子节点属性值

plotNode(firstStr, cntrPt, parentPt, decisionNode)

secondDict = myTree[firstStr] # 获取节点下的内容

plotTree.yOff = plotTree.yOff - 1.0/plotTree.totalD # 减少 y 偏移，树是自顶向下画的

for key in secondDict.keys(): # 键值：0、1

if type(secondDict[key]).__name__ == 'dict': # 判断是 dict 还是 value

plotTree(secondDict[key], cntrPt, str(key)) # 递归调用

else:

plotTree.xOff = plotTree.xOff + 1.0/plotTree.totalW # 更新 x 值

plotNode(secondDict[key], (plotTree.xOff, plotTree.yOff), cntrPt, leafNode)

plotMidText((plotTree.xOff, plotTree.yOff), cntrPt, str(key))

plotTree.yOff = plotTree.yOff + 1.0/plotTree.totalD

更新 createPlot()函数

def createPlot(inTree):

fig = plt.figure(1, facecolor='white')

fig.clf()

axprps = dict(xticks=[], yticks=[])

createPlot.ax1 = plt.subplot(111, frameon=False, **axprps) # 定义绘图区

plotTree.totalW = float(getNumLeafs(inTree)) # 存储树的宽度

plotTree.totalD = float(getTreeDepth(inTree)) #存储树的深度

# 使用了这两个全局变量追踪已经绘制的节点位置，以及放置下一个节点的恰当位置

plotTree.xOff = -0.5/plotTree.totalW;plotTree.yOff = 1.0;

plotTree(inTree, (0.5,1.0), ' ')

plt.show()

绘制图形的 x 轴有效范围是 0.0 到 1.0，y 轴有效范围也是 0.0 到 1.0。

按照图形比例绘制树形图的最大好处是无需关心实际输出图形的大小，一旦图形大小发生了变化，函数会自动按照图形大小重新绘制。

运行结果：

#变更字典，重新绘制

myTree['no surfacing'][3]='maybe'

运行结果：

END.