#coding=utf-8importmatplotlib.pyplotasplt#定义文本框和箭头格式decisionNode=dict(boxstyle="sawtooth",fc="0.8")leafNode = dict(boxstyle="round4",fc="0.8")arrow_args=dict(arrowstyle="<-")defgetNumLeafs(myTree):numLeafs=0 firstStr=myTree.keys()[0]secondDict=myTree[firstStr]forkeyinsecondDict.keys():iftype(secondDict[key]).__name__=='dict':numLeafs+=getNumLeafs(secondDict[key])else:numLeafs+=1returnnumLeafsdefgetTreeDepth(myTree):maxDepth=0 firstStr=myTree.keys()[0]secondDict = myTree[firstStr]forkeyinsecondDict.keys():iftype(secondDict[key]).__name__=='dict':thisDepth=1+getTreeDepth(secondDict[key])else:thisDepth=1ifthisDepth>maxDepth:maxDepth=thisDepthreturnmaxDepthdefplotMidText(cntrPt,parentPt,txtString):xMid=(parentPt[0]-cntrPt[0])/2.0+cntrPt[0]yMid = (parentPt[1] - cntrPt[1]) / 2.0 + cntrPt[1]createPlot.axl.text(xMid,yMid,txtString)defplotTree(myTree,parentPt,nodeTxt):numLeafs=getNumLeafs(myTree)depth=getTreeDepth(myTree)firstStr=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.totalDforkeyinsecondDict.keys():iftype(secondDict[key]).__name__=='dict':plotTree(secondDict[key],cntrPt,str(key))else:plotTree.xOff=plotTree.xOff+1.0/plotTree.totalWplotNode(secondDict[key],(plotTree.xOff,plotTree.yOff),cntrPt,leafNode)plotMidText((plotTree.xOff,plotTree.yOff),cntrPt,str(key))plotTree.yOff=plotTree.yOff+1.0/plotTree.totalDdefcreatePlot(inTree):fig=plt.figure(1,facecolor='white')fig.clf()axprops=dict(xticks=[],yticks=[])createPlot.axl=plt.subplot(111,frameon=False,**axprops)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()defplotNode(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)# def createPlot():#fig=plt.figure(1,facecolor='white')#fig.clf()#createPlot.axl=plt.subplot(111,frameon=False)#plotNode('决策节点',(0.5,0.1),(0.1,0.5),decisionNode)#plotNode('叶节点', (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'}}}}]returnlistOfTrees[i]# myTree=retrieveTree(0)# myTree['no surfacing'][3]='maybe'# print myTree# createPlot(myTree)# print retrieveTree(0)# print getNumLeafs(myTree=retrieveTree(0))# print getTreeDepth(myTree=retrieveTree(0))
deffileReading(filename):fr=open(filename)dataSet=fr.readlines()lenses=[inst.strip().split('\t')forinstindataSet]lensesLabels=['age','prescript','astigmatic','tearRate']returnlenses,lensesLabelsdataSet,labels=fileReading('lenses.txt')lensesTree=createTree(dataSet,labels)printlensesTreetreePlotter.createPlot(lensesTree)
