# -*- coding: utf-8 -*-import sys
import os
import numpy as np
import matplotlib.pyplot as plt
import treePlotter as tp
# 绘制树
myTree ={'root':{0:'leaf node',1:{'level 2':{0:'leaf node',1:'leaf node'}},2:{'level2':{0:'leaf node',1:'leaf node'}}}}
tp.createPlot(myTree)
treePlotter.py
'''
Created on Oct 14, 2010
@author: Peter Harrington
'''import matplotlib.pyplot as plt
decisionNode =dict(boxstyle="sawtooth", fc="0.8")
leafNode =dict(boxstyle="round4", fc="0.8")
arrow_args =dict(arrowstyle="<-")defgetNumLeafs(myTree):
numLeafs =0
firstStr =list(myTree.keys())[0]
secondDict = myTree[firstStr]for key in secondDict.keys():iftype(secondDict[key]).__name__=='dict':#test to see if the nodes are dictonaires, if not they are leaf nodes
numLeafs += getNumLeafs(secondDict[key])else: numLeafs +=1return numLeafs
defgetTreeDepth(myTree):
maxDepth =0
firstStr =list(myTree.keys())[0]
secondDict = myTree[firstStr]for key in secondDict.keys():iftype(secondDict[key]).__name__=='dict':#test to see if the nodes are dictonaires, if not they are leaf nodes
thisDepth =1+ getTreeDepth(secondDict[key])else: thisDepth =1if thisDepth > maxDepth: maxDepth = thisDepth
return maxDepth
defplotNode(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 )defplotMidText(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)defplotTree(myTree, parentPt, nodeTxt):#if the first key tells you what feat was split on
numLeafs = getNumLeafs(myTree)#this determines the x width of this tree
depth = getTreeDepth(myTree)
firstStr =list(myTree.keys())[0]#the text label for this node should be this
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
for key in secondDict.keys():iftype(secondDict[key]).__name__=='dict':#test to see if the nodes are dictonaires, if not they are leaf nodes
plotTree(secondDict[key],cntrPt,str(key))#recursionelse:#it's a leaf node print the leaf node
plotTree.xOff = plotTree.xOff +1.0/plotTree.totalW
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
#if you do get a dictonary you know it's a tree, and the first element will be another dictdefcreatePlot(inTree):
fig = plt.figure(1, facecolor='white')
fig.clf()
axprops =dict(xticks=[], yticks=[])
createPlot.ax1 = plt.subplot(111, frameon=False,**axprops)#no ticks#createPlot.ax1 = plt.subplot(111, frameon=False) #ticks for demo puropses
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()#def createPlot():# fig = plt.figure(1, facecolor='white')# fig.clf()# createPlot.ax1 = plt.subplot(111, frameon=False) #ticks for demo puropses # plotNode('a decision node', (0.5, 0.1), (0.1, 0.5), decisionNode)# plotNode('a leaf node', (0.8, 0.1), (0.3, 0.8), leafNode)# plt.show()defretrieveTree(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]#createPlot(thisTree)