devbox
IT小白
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

决策树代码实现

作者：IT小白 | 2024-08-22 03:49:29

赞

踩

决策树代码

目录

一、前言

二、代码实现逻辑

        构造树模块

        选择最好特征模块

        计算熵值模块

        切分数据集模块

        当前样本中多数类别模块

三、可视化拓展

 四、结果展示+完整代码

一、前言

        本文需要读者有对决策树有一定的基础，可以参考决策树原理(决策树算法概述，熵，信息增益，信息增益率，gini系数，剪枝，回归、分类任务解决)

二、代码实现逻辑

        构造树模块

        （学过数据结构的都知道，构造树最好的方法是递归）

        1.判断是否需要建树：如果当前节点所有样本的标签相同，不需要建树，如果所有特征都用完还是没有完全分类好，则分类结果采取需要少数服从多数的策略。

        2.把最好的那个特征选出来用来当作根节点

        3.根据根节点的不同特征值进行分叉

        4.在数据集中把以根节点为特征的特征值去掉（更新数据集）

        5.在特征值里循环递归建树

        6.返回树

        注意：采用字典嵌套的形式来存储树，featLabels表示根节点的值，可以根据先后顺序把特征值存储起来。

        

  1. def crecateTree(dataset,labels,featLabels):
  2. 	'''
  4. 	:param dataset: 数据集
  5. 	:param labels: 判断当前节点是否需要再分
  6. 	:param featLabels: 根节点的值
  7. 	:return:
  8. 	'''
  9. 	classList = [example[-1] for example in dataset] #当前节点的所有样本的标签
  10. 	if classList.count(classList[0]) == len(classList): #判断所有标签是否一致
  11. 		return classList[0]
  12. 	if len(dataset[0]) == 1: #只剩下一列特征值
  13. 		return majorityCnt(classList) #返回主要特征
  14. 	bestFeature = chooseBestFeatureToSplit(dataset) #得到最好特征的索引
  15. 	bestFeatureLabel = labels[bestFeature]
  16. 	featLabels.append(bestFeatureLabel)
  17. 	myTree = {bestFeatureLabel:{}} #用字典来存储树，嵌套
  18. 	del labels[bestFeature] #删除特征值
  19. 	featValue = [example[bestFeature] for example in dataset] #得到根节点特征值
  20. 	uniqueVals = set(featValue)# 有几个不同的特征值，树分几个叉
  21. 	for value in uniqueVals: #递归调用
  22. 		sublabels = labels[:]
  23. 		myTree[bestFeatureLabel][value] = crecateTree(splitDataSet(dataset,bestFeature,value),sublabels,featLabels)
  24. 	return myTree
'
运行

        选择最好特征模块

        需要把每个特征都遍历一遍，选择信息增益最大的那个特征

        

  1. def chooseBestFeatureToSplit(dataset): #核心，熵值计算
  2. 	numFeatures = len(dataset[0]) - 1 #特征的个数 随便一列减去label
  3. 	baseEntropy = calcShannonEnt(dataset) #计算当前什么都不做的熵值
  4. 	bestInfoGain = 0 #最好的信息增益
  5. 	bestFeature = -1 #最好的特征
  6. 	for i in range(numFeatures):
  7. 		featList = [example[i] for example in dataset] #当前的特征列
  8. 		uniqueVals = set(featList) #特征值的类别
  9. 		newEntropy = 0
  10. 		for val in uniqueVals:
  11. 			subDataSet = splitDataSet(dataset,i,val)
  12. 			prob = len (subDataSet) / float(len(dataset))
  13. 			newEntropy += prob * calcShannonEnt(subDataSet) # 选择特征后的熵值
  14. 		infoGain = baseEntropy - newEntropy
  15. 		if(infoGain > bestInfoGain):
  16. 			bestInfoGain = infoGain
  17. 			bestFeature = i
  18. 	return bestFeature
'
运行

        计算熵值模块

        把需要的概率算出来

  1. def calcShannonEnt(dataset):#熵值计算
  2. 	numexamples = len(dataset)
  3. 	labelCount = {}
  4. 	for featVec in dataset:
  5. 		currentlabel = featVec[-1]
  6. 		if currentlabel not in labelCount.keys():
  7. 			labelCount[currentlabel] = 0
  8. 		labelCount[currentlabel] += 1
  9.  
  10. 	shannonEnt = 0
  11. 	for key in labelCount:
  12. 		prop = float(labelCount[key]/numexamples) #概率值
  13. 		shannonEnt -= prop*log(prop,2)  #熵值
  14. 	return shannonEnt
'
运行

        切分数据集模块

        每次进行划分后都需要数据切分，包括去掉根节点特征的那一列

  1. def splitDataSet(dataset,axis,val): #切分数据集，把根节点的那一特征列去掉
  2. 	retDataSet = []
  3. 	for featVec in dataset:
  4. 		if featVec[axis] == val:
  5. 			reducedFeatVec = featVec[:axis]
  6. 			reducedFeatVec.extend(featVec[axis+1:]) #用切片和拼接把第axis列切掉
  7. 			retDataSet.append(reducedFeatVec)
  8. 	return retDataSet
'
运行

        当前样本中多数类别模块

        当所有的特征都用完后还不能完全划分，采取少数服从多数策略

  1. def majorityCnt(classList):		#当前多数类别是哪一个
  2. 	classCount = {}
  3. 	for vote in classList:
  4. 		if vote not in classCount.keys():
  5. 			classCount[vote] = 0
  6. 		classCount[vote] += 1
  7. 	sortedClassCount = sorted(classCount.items(),key=operator.itemgetter(1),reverse=True) #排序
  8. 	return sortedClassCount[0][0]
'
运行

三、可视化拓展

        这个不是重点，重要的是掌握递归建树的思想

        

  1. def getNumLeafs(myTree):
  2. 	numLeafs = 0
  3. 	firstStr = next(iter(myTree))
  4. 	secondDict = myTree[firstStr]
  5. 	for key in secondDict.keys():
  6. 		if type(secondDict[key]).__name__=='dict':
  7. 			numLeafs += getNumLeafs(secondDict[key])
  8. 		else:
  9. 			numLeafs +=1
  10. 	return numLeafs
  11.  
  12.  
  13. def getTreeDepth(myTree):
  14. 	maxDepth = 0
  15. 	firstStr = next(iter(myTree))
  16. 	secondDict = myTree[firstStr]
  17. 	for key in secondDict.keys():
  18. 		if type(secondDict[key]).__name__=='dict':
  19. 			thisDepth = 1 + getTreeDepth(secondDict[key])
  20. 		else:
  21. 			thisDepth = 1
  22. 		if thisDepth > maxDepth: maxDepth = thisDepth
  23. 	return maxDepth
  24.  
  25. def plotNode(nodeTxt, centerPt, parentPt, nodeType):
  26. 	arrow_args = dict(arrowstyle="<-")
  27. 	font = FontProperties(fname=r"c:\windows\fonts\simsunb.ttf", size=14)
  28.  
  29. 	createPlot.axl.annotate(nodeTxt, xy=parentPt, xycoords='axes fraction',
  30. 							xytext=centerPt, textcoords='axes fraction',
  31. 							va="center", ha="center", bbox=nodeType, arrowprops=arrow_args, FontProperties=font)
  32.  
  33.  
  34. def plotMidText(cntrPt, parentPt, txtString):
  35. 	xMid = (parentPt[0]-cntrPt[0])/2.0 + cntrPt[0]
  36. 	yMid = (parentPt[1]-cntrPt[1])/2.0 + cntrPt[1]
  37. 	createPlot.axl.text(xMid, yMid, txtString, va="center", ha="center", rotation=30)
  38.  
  39.  
  40. def plotTree(myTree, parentPt, nodeTxt):
  41. 	decisionNode = dict(boxstyle="sawtooth", fc="0.8")
  42. 	leafNode = dict(boxstyle="round4", fc="0.8")
  43. 	numLeafs = getNumLeafs(myTree)
  44. 	depth = getTreeDepth(myTree)
  45. 	firstStr = next(iter(myTree))
  46. 	cntrPt = (plotTree.xOff + (1.0 + float(numLeafs))/2.0/plotTree.totalW, plotTree.yOff)
  47. 	plotMidText(cntrPt, parentPt, nodeTxt)
  48. 	plotNode(firstStr, cntrPt, parentPt, decisionNode)
  49. 	secondDict = myTree[firstStr]
  50. 	plotTree.yOff = plotTree.yOff - 1.0/plotTree.totalD
  51. 	for key in secondDict.keys():
  52. 		if type(secondDict[key]).__name__=='dict':
  53. 			plotTree(secondDict[key],cntrPt,str(key))
  54. 		else:
  55. 			plotTree.xOff = plotTree.xOff + 1.0/plotTree.totalW
  56. 			plotNode(secondDict[key], (plotTree.xOff, plotTree.yOff), cntrPt, leafNode)
  57. 			plotMidText((plotTree.xOff, plotTree.yOff), cntrPt, str(key))
  58. 	plotTree.yOff = plotTree.yOff + 1.0/plotTree.totalD
  59.  
  60.  
  61. def createPlot(inTree):
  62. 	fig = plt.figure(1, facecolor='white')													#创建fig
  63. 	fig.clf()																				#清空fig
  64. 	axprops = dict(xticks=[], yticks=[])
  65. 	createPlot.ax1 = plt.subplot(111, frameon=False, **axprops)    							#去掉x、y轴
  66. 	plotTree.totalW = float(getNumLeafs(inTree))											#获取决策树叶结点数目
  67. 	plotTree.totalD = float(getTreeDepth(inTree))											#获取决策树层数
  68. 	plotTree.xOff = -0.5/plotTree.totalW; plotTree.yOff = 1.0							#x偏移
  69. 	plotTree(inTree, (0.5,1.0), '')															#绘制决策树
  70. 	plt.show()
'
运行

 四、结果展示+完整代码

  1. # -*- coding: UTF-8 -*-
  2. from matplotlib.font_manager import FontProperties
  3. import matplotlib.pyplot as plt
  4. from math import log
  5. import operator
  6.  
  7.  
  8.  
  9. def createDataSet():
  10. 	dataSet = [[0, 0, 0, 0, 'no'],						
  11. 			[0, 0, 0, 1, 'no'],
  12. 			[0, 1, 0, 1, 'yes'],
  13. 			[0, 1, 1, 0, 'yes'],
  14. 			[0, 0, 0, 0, 'no'],
  15. 			[1, 0, 0, 0, 'no'],
  16. 			[1, 0, 0, 1, 'no'],
  17. 			[1, 1, 1, 1, 'yes'],
  18. 			[1, 0, 1, 2, 'yes'],
  19. 			[1, 0, 1, 2, 'yes'],
  20. 			[2, 0, 1, 2, 'yes'],
  21. 			[2, 0, 1, 1, 'yes'],
  22. 			[2, 1, 0, 1, 'yes'],
  23. 			[2, 1, 0, 2, 'yes'],
  24. 			[2, 0, 0, 0, 'no']]
  25. 	labels = ['F1-AGE', 'F2-WORK', 'F3-HOME', 'F4-LOAN']		
  26. 	return dataSet, labels
  27.  
  28. def crecateTree(dataset,labels,featLabels):
  29. 	'''
  31. 	:param dataset: 数据集
  32. 	:param labels: 判断当前节点是否需要再分
  33. 	:param featLabels: 节点的值
  34. 	:return:
  35. 	'''
  36. 	classList = [example[-1] for example in dataset] #当前节点的所有样本的标签
  37. 	if classList.count(classList[0]) == len(classList): #判断所有标签是否一致
  38. 		return classList[0]
  39. 	if len(dataset[0]) == 1: #只剩下一列特征值
  40. 		return majorityCnt(classList) #返回主要特征
  41. 	bestFeature = chooseBestFeatureToSplit(dataset) #得到最好特征的索引
  42. 	bestFeatureLabel = labels[bestFeature]
  43. 	featLabels.append(bestFeatureLabel)
  44. 	myTree = {bestFeatureLabel:{}} #用字典来存储树，嵌套
  45. 	del labels[bestFeature] #删除特征值
  46. 	featValue = [example[bestFeature] for example in dataset] #得到根节点特征值
  47. 	uniqueVals = set(featValue)# 有几个不同的特征值，树分几个叉
  48. 	for value in uniqueVals: #递归调用
  49. 		sublabels = labels[:]
  50. 		myTree[bestFeatureLabel][value] = crecateTree(splitDataSet(dataset,bestFeature,value),sublabels,featLabels)
  51. 	return myTree
  52.  
  53. def majorityCnt(classList):		#当前多数类别是哪一个
  54. 	classCount = {}
  55. 	for vote in classList:
  56. 		if vote not in classCount.keys():
  57. 			classCount[vote] = 0
  58. 		classCount[vote] += 1
  59. 	sortedClassCount = sorted(classCount.items(),key=operator.itemgetter(1),reverse=True) #排序
  60. 	return sortedClassCount[0][0]
  61.  
  62. def chooseBestFeatureToSplit(dataset): #核心，熵值计算
  63. 	numFeatures = len(dataset[0]) - 1 #特征的个数 随便一列减去label
  64. 	baseEntropy = calcShannonEnt(dataset) #计算当前什么都不做的熵值
  65. 	bestInfoGain = 0 #最好的信息增益
  66. 	bestFeature = -1 #最好的特征
  67. 	for i in range(numFeatures):
  68. 		featList = [example[i] for example in dataset] #当前的特征列
  69. 		uniqueVals = set(featList) #特征值的类别
  70. 		newEntropy = 0
  71. 		for val in uniqueVals:
  72. 			subDataSet = splitDataSet(dataset,i,val)
  73. 			prob = len (subDataSet) / float(len(dataset))
  74. 			newEntropy += prob * calcShannonEnt(subDataSet) # 选择特征后的熵值
  75. 		infoGain = baseEntropy - newEntropy
  76. 		if(infoGain > bestInfoGain):
  77. 			bestInfoGain = infoGain
  78. 			bestFeature = i
  79. 	return bestFeature
  80.  
  81.  
  82. def splitDataSet(dataset,axis,val): #切分数据集，把根节点的那一特征列去掉
  83. 	retDataSet = []
  84. 	for featVec in dataset:
  85. 		if featVec[axis] == val:
  86. 			reducedFeatVec = featVec[:axis]
  87. 			reducedFeatVec.extend(featVec[axis+1:]) #用切片和拼接把第axis列切掉
  88. 			retDataSet.append(reducedFeatVec)
  89. 	return retDataSet
  90.  
  91. def calcShannonEnt(dataset):#熵值计算
  92. 	numexamples = len(dataset)
  93. 	labelCount = {}
  94. 	for featVec in dataset:
  95. 		currentlabel = featVec[-1]
  96. 		if currentlabel not in labelCount.keys():
  97. 			labelCount[currentlabel] = 0
  98. 		labelCount[currentlabel] += 1
  99.  
  100. 	shannonEnt = 0
  101. 	for key in labelCount:
  102. 		prop = float(labelCount[key]/numexamples) #概率值
  103. 		shannonEnt -= prop*log(prop,2)  #熵值
  104. 	return shannonEnt
  105.  
  106.  
  107. def getNumLeafs(myTree):
  108. 	numLeafs = 0
  109. 	firstStr = next(iter(myTree))
  110. 	secondDict = myTree[firstStr]
  111. 	for key in secondDict.keys():
  112. 		if type(secondDict[key]).__name__=='dict':
  113. 			numLeafs += getNumLeafs(secondDict[key])
  114. 		else:
  115. 			numLeafs +=1
  116. 	return numLeafs
  117.  
  118.  
  119. def getTreeDepth(myTree):
  120. 	maxDepth = 0
  121. 	firstStr = next(iter(myTree))
  122. 	secondDict = myTree[firstStr]
  123. 	for key in secondDict.keys():
  124. 		if type(secondDict[key]).__name__=='dict':
  125. 			thisDepth = 1 + getTreeDepth(secondDict[key])
  126. 		else:
  127. 			thisDepth = 1
  128. 		if thisDepth > maxDepth: maxDepth = thisDepth
  129. 	return maxDepth
  130.  
  131. def plotNode(nodeTxt, centerPt, parentPt, nodeType):
  132. 	arrow_args = dict(arrowstyle="<-")
  133. 	font = FontProperties(fname=r"c:\windows\fonts\simsunb.ttf", size=14)
  134.  
  135. 	createPlot.axl.annotate(nodeTxt, xy=parentPt, xycoords='axes fraction',
  136. 							xytext=centerPt, textcoords='axes fraction',
  137. 							va="center", ha="center", bbox=nodeType, arrowprops=arrow_args, FontProperties=font)
  138.  
  139.  
  140. def plotMidText(cntrPt, parentPt, txtString):
  141. 	xMid = (parentPt[0]-cntrPt[0])/2.0 + cntrPt[0]
  142. 	yMid = (parentPt[1]-cntrPt[1])/2.0 + cntrPt[1]
  143. 	createPlot.axl.text(xMid, yMid, txtString, va="center", ha="center", rotation=30)
  144.  
  145.  
  146. def plotTree(myTree, parentPt, nodeTxt):
  147. 	decisionNode = dict(boxstyle="sawtooth", fc="0.8")
  148. 	leafNode = dict(boxstyle="round4", fc="0.8")
  149. 	numLeafs = getNumLeafs(myTree)
  150. 	depth = getTreeDepth(myTree)
  151. 	firstStr = next(iter(myTree))
  152. 	cntrPt = (plotTree.xOff + (1.0 + float(numLeafs))/2.0/plotTree.totalW, plotTree.yOff)
  153. 	plotMidText(cntrPt, parentPt, nodeTxt)
  154. 	plotNode(firstStr, cntrPt, parentPt, decisionNode)
  155. 	secondDict = myTree[firstStr]
  156. 	plotTree.yOff = plotTree.yOff - 1.0/plotTree.totalD
  157. 	for key in secondDict.keys():
  158. 		if type(secondDict[key]).__name__=='dict':
  159. 			plotTree(secondDict[key],cntrPt,str(key))
  160. 		else:
  161. 			plotTree.xOff = plotTree.xOff + 1.0/plotTree.totalW
  162. 			plotNode(secondDict[key], (plotTree.xOff, plotTree.yOff), cntrPt, leafNode)
  163. 			plotMidText((plotTree.xOff, plotTree.yOff), cntrPt, str(key))
  164. 	plotTree.yOff = plotTree.yOff + 1.0/plotTree.totalD
  165.  
  166.  
  167. def createPlot(inTree):
  168. 	fig = plt.figure(1, facecolor='white')													#创建fig
  169. 	fig.clf()																				#清空fig
  170. 	axprops = dict(xticks=[], yticks=[])
  171. 	createPlot.ax1 = plt.subplot(111, frameon=False, **axprops)    							#去掉x、y轴
  172. 	plotTree.totalW = float(getNumLeafs(inTree))											#获取决策树叶结点数目
  173. 	plotTree.totalD = float(getTreeDepth(inTree))											#获取决策树层数
  174. 	plotTree.xOff = -0.5/plotTree.totalW; plotTree.yOff = 1.0							#x偏移
  175. 	plotTree(inTree, (0.5,1.0), '')															#绘制决策树
  176. 	plt.show()
  177.  
  178. if __name__ == '__main__':
  179. 	dataSet,labels = createDataSet()
  180. 	featLabels = []
  181. 	myTree = crecateTree(dataSet, labels, featLabels)
  182. 	print(featLabels)
  183. 	createPlot(myTree)

 

        选择两个特征建树

        可视化结果：

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/IT小白/article/detail/1014588
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号