机器学习实战-决策树-画图

 

获取树节点的数据以及树的层数

1. def getNumLeafs(myTree):  
2.     numLeafs=0  
3. firstStr=list(myTree.keys())[0]   
4. #此处注意，书上用的是py2的版本，所以有些不同，对于a.keys()的类型需要先转成列表  
5. secondDict=myTree[firstStr]  
6. #firstStr取出决策树的顶级父节点，然后读取顶级父节点所建立的树  
7.     for key in secondDict.keys():  
8.         if type(secondDict[key]).__name__=='dict': #判断某个键值是否是一个dict类型，即是否是一个决策节点  
9.              numLeafs+=getNumLeafs(secondDict[key])#递归算法，不断递归，直至找到叶子节点  
10.         else:  
11.             numLeafs+=1  
12.     return numLeafs  
13.     
14. def getTreeDepth(myTree):  
15.     maxDepth=0  
16.     firstStr=list(myTree.keys())[0]  
17.     secondDict=myTree[firstStr]  
18.     for key in secondDict.keys():  
19.         if type(secondDict[key]).__name__=='dict':  
20. #上面内容都一样  
21.                 hisDepth=1+getTreeDepth(secondDict[key])#如果是决策节点，说明就需要至少还有一层，所以加上1，并搜寻此决策节点下是否还有决策节点  
22.         else:  
23.             thisDepth=1  
24.         if thisDepth>maxDepth:  
25.             maxDepth=thisDepth  
26. return maxDepth