机器学习决策树算法之代码实现

上篇文章讨论了决策树算法的原理，这一节直接附上代码，代码整理自博友，链接如下：

https://blog.csdn.net/c406495762/article/details/76262487

# 年龄：0代表青年，1代表中年，2代表老年；
# 有工作：0代表否，1代表是；
# 有自己的房子：0代表否，1代表是；
# 信贷情况：0代表一般，1代表好，2代表非常好；
# 类别(是否给贷款)：no代表否，yes代表是
from math import log
import operator
from matplotlib.font_manager import FontProperties
import matplotlib.pyplot as plt
#创建数据集
def 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 = ['年龄', '有工作', '有自己的房子', '信贷情况']
    return dataSet, labels


# 计算给定数据的香农熵
def calcShannonEnt(dataSet):
    numEntires = len(dataSet)  # 返回数据集的行数
    labelCounts = {}  # 保存每个标签(Label)出现次数的字典
    for featVec in dataSet:  # 对每组特征向量进行统计
        currentLabel = featVec[-1]  # 提取标签(Label)信息（yes or no）
        if currentLabel not in labelCounts.keys():  # 如果没有当前的标签，添加该标签，
            labelCounts[currentLabel] = 0
        labelCounts[currentLabel] += 1  # Label计数，返回的是一个字典{'no':6,'yes':9}
    shannonEnt = 0.0  # 经验熵(香农熵)
    for key in labelCounts:  # 计算香农熵
        prob = float(labelCounts[key]) / numEntires  # 选择该标签(Label)的概率
        shannonEnt -= prob * log(prob, 2)  # 利用公式计算
    return shannonEnt  # 返回经验熵(香农熵)


#按照给定特征划分数据集
#axis表示数据集中的特征项,（即该数据中的列标签的索引）该数据集中特征项是['年龄', '有工作', '有自己的房子', '信贷情况']
#value表示某一个特征项中的某一种情况，比如有自己的房子就是{0,1,}，分别代表无房子，有房子
def splitDataSet(dataSet,axis,value):
    retDataSet=[]    #先设置一个空数组
    for featVec in dataSet:#对dataSet的每一行进行遍历
        #取一行的第axis个值（这个值表示特征项），与value值进行比较
        if featVec[axis] == value:
            reduceFeatVec = featVec[:axis]#这里的reduceFeatVec是一个空数组
            reduceFeatVec.extend(featVec[axis+1:])#这里的reduceFeatVec表示除开特征项的其他项
            retDataSet.append(reduceFeatVec)#将剔除特征项的数组全部相加组成一个新的矩阵
    return retDataSet




    # [['0', '0', '0', 'no'],
    #  ['0', '0', '1', 'no'],
    #  ['0', '1', '1', 'yes'],
    #  ['0', '1', '0', 'yes'],
    #  ['0', '0', '0', 'no'], 
    #  ['1', '0', '0', 'no'],
    #  ['1', '0', '1', 'no'],
    #  ['1', '1', '1', 'yes'],
    
    
    #  ['1', '0', '2', 'yes'],
    #  ['1', '0', '2', 'yes'],
    #  ['2', '0', '2', 'yes'],
    #  ['2', '0', '1', 'yes'],
    #  ['2', '1', '1', 'yes'],
    #  ['2', '1', '2', 'yes'],
    #  ['2', '0', '0', 'no']],


##选择给定数据集中的最优特征
def chooseBestFeatureToSplit(dataSet):
    numFeatures = len(dataSet[0])-1#获取到特征的数量，在这里返回的是4个特征：4
    baseEntropy = calcShannonEnt(dataSet)#首先计算香农熵
    bestInfoGain = 0.0#定义信息增益
    bestFeature = -1  #定义最优特征索引值
    for i in range(numFeatures):
        featList = [example[i] for example in dataSet]#这个list将获取到某一个特征项的所有取值：比如第一个特征项的取值：包括了15个0和1
        uniqueVals = set(featList)#将获取到的所有情况去重
        newEntropy = 0.0#定义信息增益
        for value in uniqueVals:
            subDataSet = splitDataSet(dataSet, i, value)#subDataSet表示每一个value取值下的数据集
            prob = len(subDataSet) / float(len(dataSet))#每种情况所占的比例
            newEntropy += prob * calcShannonEnt(subDataSet)#求经验熵
        infoGain = baseEntropy - newEntropy#计算出该特征下的信息增益
        if (infoGain >bestInfoGain):#比较每个特征的信息增益
            bestInfoGain = infoGain#选取信息增益最大的值
            bestFeature = i#该索引对应的特征即为信息增大的特征项
    return bestFeature
    

#统计classList中出现此处最多的元素(类标签)
#classList是某个类标签的所有情况，
# 比如：['no', 'no', 'yes', 'yes', 'no', 'no', 'no', 'yes', 'yes', 'yes', 'yes', 'yes', 'yes', 'yes', 'no']
def majorityCnt(classList):
    classCount = {}                                              #创建一个空字典
    for vote in classList:                                       #将classList中每个元素进行遍历
        if vote not in classCount.keys():classCount[vote] = 0    #如果该元素不在字典的key里，则在字典里创建该元素
        classCount[vote] += 1                                    #如果该元素在字典key里，则对元素值+1，最后得到的结果类似于{'no':6,'yes':9}
    sortedClassCount = sorted(classCount.items(), key = operator.itemgetter(1), reverse = True)#根据字典的值降序排序,结果类似于{'yes':9,'no':6}
    return sortedClassCount[0][0]                                #返回classList中出现次数最多的元素,比如'yes'



#函数说明:创建决策树
def createTree(dataSet, labels, featLabels):
    classList = [example[-1] for example in dataSet]               #取分类标签(是否放贷:yes or no)
    if classList.count(classList[0]) == len(classList):            #如果classList的值全都是一样的，
        return classList[0]                                        #则返回classList的第一个元素(第一个元素也是整个list的元素)
    if len(dataSet[0]) == 1:                                       #遍历完所有特征时返回出现次数最多的类标签
        return majorityCnt(classList)
    bestFeat = chooseBestFeatureToSplit(dataSet)                   #选择最优特征的索引
    bestFeatLabel = labels[bestFeat]                               #最优特征索引对应的标签
    featLabels.append(bestFeatLabel)                               #featLabels里增加一个标签元素['有房子']
    myTree = {bestFeatLabel:{}}                                    #根据最优特征的标签生成树 {'有房子': {}}
    del(labels[bestFeat])                                          #删除已经使用特征标签 ，剩下的就是['年龄', '有工作', '信贷情况']
    featValues = [example[bestFeat] for example in dataSet]        #得到训练集中所有最优特征的属性值，[0, 0, 0, 1, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0]
    uniqueVals = set(featValues)                                   #去掉重复的属性值{0, 1}
    for value in uniqueVals:                                       #遍历特征，创建决策树。
        #递归调用createTree函数，如果分类过程中已经是叶子节点了，直接返回叶子结点的值，
        myTree[bestFeatLabel][value] = createTree(splitDataSet(dataSet, bestFeat, value), labels, featLabels)
    return myTree


#使用决策树分类
def classify(inputTree, featLabels, testVec):                            #featLabels=['有自己的房子', '有工作']
    firstStr = next(iter(inputTree))                                     #获取决策树第一个节点的标签，也就是根节点，在这里是'有自己的房子'
    secondDict = inputTree[firstStr]                                     #根节点对应的value值，secondDict={0: {'有工作': {0: 'no', 1: 'yes'}}, 1: 'yes'}
    featIndex = featLabels.index(firstStr)                               #取根节点的索引
    for key in secondDict.keys():                                        #对secondDict的key进行遍历
        if testVec[featIndex] == key:                                    #当遍历到key与测试数据的第一个特征值相等时
            if type(secondDict[key]).__name__ == 'dict':                 #如果测试数据的特征值在决策树里对应的不是叶子节点，
                classLabel = classify(secondDict[key], featLabels, testVec) #则递归是引用该分类方法，知道能返回叶子节点
            else: classLabel = secondDict[key]                           #如果正好是叶子节点，直接返回结果
    return classLabel



#决策树的叶子结点的数目
def getNumLeafs(myTree):
    numLeafs = 0
    firstStr = next(iter(myTree))     #firstStr返回第一个分类标签：有自己的房子
    secondDict = myTree[firstStr]     #secondDict返回的是第二层决策树分类：{0: {'有工作': {0: 'no', 1: 'yes'}}, 1: 'yes'}
    for key in secondDict.keys():    
        if type(secondDict[key]).__name__=='dict':
              numLeafs += getNumLeafs(secondDict[key])#递归操作，分完一次类，找到一个叶子节点就加1
        else:   numLeafs +=1
    return numLeafs



#获取决策树的层数
def getTreeDepth(myTree):
    maxDepth = 0                                    
    firstStr = next(iter(myTree))       #获取第一个特征项，firstStr='有自己的房子'
    secondDict = myTree[firstStr]       #获取下一个节点，secondDict={0: {'有工作': {0: 'no', 1: 'yes'}}, 1: 'yes'}
    for key in secondDict.keys():
        if type(secondDict[key]).__name__=='dict':                #测试该结点是否为字典，
            thisDepth = 1 + getTreeDepth(secondDict[key])         #是字典，则继续递归调用函数
        else:   thisDepth = 1                                     #否则是叶子节点，层数加1
        if thisDepth > maxDepth: maxDepth = thisDepth             #更新层数
    return maxDepth


#绘制节点
def plotNode(nodeTxt, centerPt, parentPt, nodeType):
    arrow_args = dict(arrowstyle="<-")                                            #定义箭头格式
    font = FontProperties(fname=r"c:\windows\fonts\simsun.ttc", size=14)          #设置中文字体
    createPlot.ax1.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.ax1.text(xMid, yMid, txtString, va="center", ha="center", rotation=30)
    
#绘制决策树
def plotTree(myTree, parentPt, nodeTxt):
    decisionNode = dict(boxstyle="sawtooth", fc="0.8")                                        #设置结点格式
    leafNode = dict(boxstyle="round4", 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.totalD                                       #y偏移
    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
    
    
    
#创建绘制面板
def createPlot(inTree):
    fig = plt.figure(1, facecolor='white')                                                    #创建fig
    fig.clf()                                                                                 #清空fig
    axprops = 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 = createTree(dataSet, labels, featLabels)
    print('决策树结构是:',myTree)
    print('作为分类的标签是：',featLabels)
    testVec = [0, 1]  # 测试数据，第一个0代表无房，第二个1代表有工作
    result = classify(myTree, featLabels, testVec)
    if result == 'yes':
        print('放贷')
    if result == 'no':
        print('不放贷')
    createPlot(myTree)

 

最后的输出结果是：

决策树结构是: {'有自己的房子': {0: {'有工作': {0: 'no', 1: 'yes'}}, 1: 'yes'}}
作为分类的标签是： ['有自己的房子', '有工作']
放贷

决策树的图形如下：