决策树代码实现

���录

一、前言

二、代码实现逻辑

        构造树模块

        选择最好特征模块

        计算熵值模块

        切分数据集模块

        当前样本中多数类别模块

三、可视化拓展

 四、结果展示+完整代码

---

一、前言

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

二、代码实现逻辑

        构造树模块

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

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

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

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

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

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

        6.返回树

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

        

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

        选择最好特征模块

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

        

def chooseBestFeatureToSplit(dataset): #核心，熵值计算numFeatures = len(dataset[0]) - 1 #特征的个数 随便一列减去labelbaseEntropy = calcShannonEnt(dataset) #计算当前什么都不做的熵值bestInfoGain = 0 #最好的信息增益bestFeature = -1 #最好的特征for i in range(numFeatures):featList = [example[i] for example in dataset] #当前的特征列uniqueVals = set(featList) #特征值的类别newEntropy = 0for val in uniqueVals:subDataSet = splitDataSet(dataset,i,val)prob = len (subDataSet) / float(len(dataset))newEntropy += prob * calcShannonEnt(subDataSet) # 选择特征后的熵值infoGain = baseEntropy - newEntropyif(infoGain > bestInfoGain):bestInfoGain = infoGainbestFeature = ireturn bestFeature

        计算熵值模块

        把需要的概率算出来

def calcShannonEnt(dataset):#熵值计算numexamples = len(dataset)labelCount = {}for featVec in dataset:currentlabel = featVec[-1]if currentlabel not in labelCount.keys():labelCount[currentlabel] = 0labelCount[currentlabel] += 1shannonEnt = 0for key in labelCount:prop = float(labelCount[key]/numexamples) #概率值shannonEnt -= prop*log(prop,2)  #熵值return shannonEnt

        切分数据集模块

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

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

        当前样本中多数类别模块

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

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

三、可视化拓展

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

        

def getNumLeafs(myTree):numLeafs = 0firstStr = next(iter(myTree))secondDict = myTree[firstStr]for key in secondDict.keys():if type(secondDict[key]).__name__=='dict':numLeafs += getNumLeafs(secondDict[key])else:numLeafs +=1return numLeafsdef getTreeDepth(myTree):maxDepth = 0firstStr = next(iter(myTree))secondDict = myTree[firstStr]for key in secondDict.keys():if type(secondDict[key]).__name__=='dict':thisDepth = 1 + getTreeDepth(secondDict[key])else:thisDepth = 1if thisDepth > maxDepth: maxDepth = thisDepthreturn maxDepthdef plotNode(nodeTxt, centerPt, parentPt, nodeType):arrow_args = dict(arrow)font = FontProperties(fname=r"c:\windows\fonts\simsunb.ttf", size=14)createPlot.axl.annotate(nodeTxt, xy=parentPt, xycoords='axes fraction',xytext=centerPt, textcoords='axes fraction',va="center", ha="center", bbox=nodeType, arrowprops=arrow_args, FontProperties=font)def plotMidText(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, va="center", ha="center", rotation=30)def plotTree(myTree, parentPt, nodeTxt):decisionNode = dict(box, fc="0.8")leafNode = dict(box, fc="0.8")numLeafs = getNumLeafs(myTree)depth = getTreeDepth(myTree)firstStr = next(iter(myTree))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.totalDfor key in secondDict.keys():if type(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.totalDdef createPlot(inTree):fig = plt.figure(1, facecolor='white')#创建figfig.clf()#清空figaxprops = dict(xticks=[], yticks=[])createPlot.ax1 = plt.subplot(111, frameon=False, **axprops)    #去掉x、y轴plotTree.totalW = float(getNumLeafs(inTree))#获取决策树叶结点数目plotTree.totalD = float(getTreeDepth(inTree))#获取决策树层数plotTree.xOff = -0.5/plotTree.totalW; plotTree.yOff = 1.0#x偏移plotTree(inTree, (0.5,1.0), '')#绘制决策树plt.show()

 四、结果展示+完整代码

# -*- coding: UTF-8 -*-from matplotlib.font_manager import FontPropertiesimport matplotlib.pyplot as pltfrom math import logimport operatordef createDataSet():dataSet = [[0, 0, 0, 0, 'no'],[0, 0, 0, 1, 'no'],[0, 1, 0, 1, 'yes'],[0, 1, 1, 0, 'yes'],[0, 0, 0, 0, 'no'],[1, 0, 0, 0, 'no'],[1, 0, 0, 1, 'no'],[1, 1, 1, 1, 'yes'],[1, 0, 1, 2, 'yes'],[1, 0, 1, 2, 'yes'],[2, 0, 1, 2, 'yes'],[2, 0, 1, 1, 'yes'],[2, 1, 0, 1, 'yes'],[2, 1, 0, 2, 'yes'],[2, 0, 0, 0, 'no']]labels = ['F1-AGE', 'F2-WORK', 'F3-HOME', 'F4-LOAN']return dataSet, labelsdef crecateTree(dataset,labels,featLabels):''':param dataset: 数据集:param labels: 判断当前节点是否需要再分:param featLabels: 节点的值:return:'''classList = [example[-1] for example in dataset] #当前节点的所有样本的标签if classList.count(classList[0]) == len(classList): #判断所有标签是否一致return classList[0]if len(dataset[0]) == 1: #只剩下一列特征值return majorityCnt(classList) #返回主要特征bestFeature = chooseBestFeatureToSplit(dataset) #得到最好特征的索引bestFeatureLabel = labels[bestFeature]featLabels.append(bestFeatureLabel)myTree = {bestFeatureLabel:{}} #用字典来存储树，嵌套del labels[bestFeature] #删除特征值featValue = [example[bestFeature] for example in dataset] #得到根节点特征值uniqueVals = set(featValue)# 有几个不同的特征值，树分几个叉for value in uniqueVals: #递归调用sublabels = labels[:]myTree[bestFeatureLabel][value] = crecateTree(splitDataSet(dataset,bestFeature,value),sublabels,featLabels)return myTreedef majorityCnt(classList):#当前多数类别是哪一个classCount = {}for vote in classList:if vote not in classCount.keys():classCount[vote] = 0classCount[vote] += 1sortedClassCount = sorted(classCount.items(),key=operator.itemgetter(1),reverse=True) #排序return sortedClassCount[0][0]def chooseBestFeatureToSplit(dataset): #核心，熵值计算numFeatures = len(dataset[0]) - 1 #特征的个数 随便一列减去labelbaseEntropy = calcShannonEnt(dataset) #计算当前什么都不做的熵值bestInfoGain = 0 #最好的信息增益bestFeature = -1 #最好的特征for i in range(numFeatures):featList = [example[i] for example in dataset] #当前的特征列uniqueVals = set(featList) #特征值的类别newEntropy = 0for val in uniqueVals:subDataSet = splitDataSet(dataset,i,val)prob = len (subDataSet) / float(len(dataset))newEntropy += prob * calcShannonEnt(subDataSet) # 选择特征后的熵值infoGain = baseEntropy - newEntropyif(infoGain > bestInfoGain):bestInfoGain = infoGainbestFeature = ireturn bestFeaturedef splitDataSet(dataset,axis,val): #切分数据集，把根节点的那一特征列去掉retDataSet = []for featVec in dataset:if featVec[axis] == val:reducedFeatVec = featVec[:axis]reducedFeatVec.extend(featVec[axis+1:]) #用切片和拼接把第axis列切掉retDataSet.append(reducedFeatVec)return retDataSetdef calcShannonEnt(dataset):#熵值计算numexamples = len(dataset)labelCount = {}for featVec in dataset:currentlabel = featVec[-1]if currentlabel not in labelCount.keys():labelCount[currentlabel] = 0labelCount[currentlabel] += 1shannonEnt = 0for key in labelCount:prop = float(labelCount[key]/numexamples) #概率值shannonEnt -= prop*log(prop,2)  #熵值return shannonEntdef getNumLeafs(myTree):numLeafs = 0firstStr = next(iter(myTree))secondDict = myTree[firstStr]for key in secondDict.keys():if type(secondDict[key]).__name__=='dict':numLeafs += getNumLeafs(secondDict[key])else:numLeafs +=1return numLeafsdef getTreeDepth(myTree):maxDepth = 0firstStr = next(iter(myTree))secondDict = myTree[firstStr]for key in secondDict.keys():if type(secondDict[key]).__name__=='dict':thisDepth = 1 + getTreeDepth(secondDict[key])else:thisDepth = 1if thisDepth > maxDepth: maxDepth = thisDepthreturn maxDepthdef plotNode(nodeTxt, centerPt, parentPt, nodeType):arrow_args = dict(arrow)font = FontProperties(fname=r"c:\windows\fonts\simsunb.ttf", size=14)createPlot.axl.annotate(nodeTxt, xy=parentPt, xycoords='axes fraction',xytext=centerPt, textcoords='axes fraction',va="center", ha="center", bbox=nodeType, arrowprops=arrow_args, FontProperties=font)def plotMidText(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, va="center", ha="center", rotation=30)def plotTree(myTree, parentPt, nodeTxt):decisionNode = dict(box, fc="0.8")leafNode = dict(box, fc="0.8")numLeafs = getNumLeafs(myTree)depth = getTreeDepth(myTree)firstStr = next(iter(myTree))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.totalDfor key in secondDict.keys():if type(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.totalDdef createPlot(inTree):fig = plt.figure(1, facecolor='white')#创建figfig.clf()#清空figaxprops = dict(xticks=[], yticks=[])createPlot.ax1 = plt.subplot(111, frameon=False, **axprops)    #去掉x、y轴plotTree.totalW = float(getNumLeafs(inTree))#获取决策树叶结点数目plotTree.totalD = float(getTreeDepth(inTree))#获取决策树层数plotTree.xOff = -0.5/plotTree.totalW; plotTree.yOff = 1.0#x偏移plotTree(inTree, (0.5,1.0), '')#绘制决策树plt.show()if __name__ == '__main__':dataSet,labels = createDataSet()featLabels = []myTree = crecateTree(dataSet, labels, featLabels)print(featLabels)createPlot(myTree)

 

        选择两个特征建树

        可视化结果：