python math算法库 决策树生成图片

2019独角兽企业重金招聘Python工程师标准>>>

每周一搏，提升自我。

前两篇文章主要介绍了基于math的ID3和C45算法决策树，这次将决策树生成图片。

#!/usr/bin/python
#coding:utf-8
'''
窗体显示决策树
'''
import matplotlib.pyplot as plt

#决策树属性设置
decisionNode=dict(boxstyle="sawtooth",fc="0.8")
leafNode=dict(boxstyle="round4",fc="0.8")
arrow_args=dict(arrowstyle="<-")


#createPlot 主函数，调用即可画出决策树，其中调用登了剩下的所有的函数，inTree的形式必须为嵌套的决策树
def createPlot(inThree):
	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(inThree))
	plotTree.totalD=float(getTreeDepth(inThree))
	plotTree.xOff=-0.5/plotTree.totalW;
	plotTree.yOff=1.0
	plotTree(inThree,(0.5,1.0),'')
	plt.show()

#决策树上节点之间的箭头设置
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 getNumLeafs(myTree):
	numLeafs=0
	firstStr=myTree.keys()[0]
	secondDict=myTree[firstStr]
	for key in secondDict.keys():
		if type(secondDict[key]).__name__=='dict':#test to see if the nodes are dictonaires, if not they are leaf nodes
			numLeafs += getNumLeafs(secondDict[key])
		else: numLeafs+=1
	return numLeafs

#得到决策树的深度
def getTreeDepth(myTree):
	maxDepthh=0
	firstStr=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>maxDepthh:maxDepthh=thisDepth
	return maxDepthh

#父子节点之间画决策树
def plotTree(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.totalD
	for key in secondDict.keys():
		if type(secondDict[key]).__name__=='dict':
			plotTree(secondDict[key],cntrPt,str(key))
		else:
			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

 

转载于:https://my.oschina.net/wangzonghui/blog/1618687