机器学习实战 决策树 代码详解

from math import log

def calcShannonEnt(dataSet):
    # 对信息熵的计算
    numEntries = len(dataSet)
    labelCounts={}
    for featVec in dataSet:
        currentLabel=featVec[-1]
        if currentLabel not in labelCounts.keys():
            labelCounts[currentLabel]=0
        labelCounts[currentLabel]+=1
    shannonEnt=0.0
    for key in labelCounts:
        prob=float(labelCounts[key])/numEntries
        shannonEnt-=prob*log(prob,2)
    return shannonEnt
def createDataSet():
    dataSet=[[1,1,'yes'],
             [1,1,'yes'],
             [1,0,'no'],
             [0,1,'no'],
             [0,1,'no']]
    lables=['no surfacing','flippers']
    return dataSet,lables
#测试数据
# myDat,labels=createDataSet()
# # '这里是把list集合看做一个字符串，-1代表的是字符串的最后一个位置
# myDat[0][-1]='maybe'
# print(myDat)
# print(calcShannonEnt(myDat))


def splitDataSet(dataSet,axis,value):
# 按照给定特征划分数据集
#1,0代表的属性的不同结果，通过splitDataSet函数 对他进行分类 例如 axis=0,value=1 就是找第0列属性中结果为1的 相同的 axis=0,value =0
#就是找第0列属性中，结果为0的
    retDataSet=[]
    for featVec in dataSet:
        if featVec[axis]==value:
            reducedFeatVec=featVec[:axis]#如果不能理解 去查找列表切片的相关知识
            reducedFeatVec.extend(featVec[axis+1:])
            retDataSet.append(reducedFeatVec)
    return retDataSet
#测试数据
# myDat,labels=createDataSet()
# print(myDat)
# print(splitDataSet(myDat,0,1))
# print(splitDataSet(myDat,0,0))

def chooseBestFeatureToSplit(dataSet):
    #选择最好的数据集划分方式
    numFeatures=len(dataSet[0])-1#数据集特征值的数目
    baseEntropy=calcShannonEnt(dataSet)#算出原始样本的信息熵
    bestInfoGain=0.0;bestFeature=-1
    for i in range(numFeatures):
        featList=[example[i] for example in dataSet]#列表解析
        #example是dataSet的每一行元素 example[i]取的是第i个索引的值
        uniqueVals=set(featList)#删除重复数据
        #uniqueVals代表的是第i个属性的可能取值
        newEntropy=0.0
        for value in uniqueVals:
            subDataSet=splitDataSet(dataSet,i,value)
            #subDataSet是第i个属性，值等于value的 分类找出的数据集合
            prob=len(subDataSet)/float(len(dataSet))
            newEntropy+=prob*calcShannonEnt(subDataSet)
            #newEntropy是分类后信息熵的计算公式
        infoGain=baseEntropy-newEntropy#信息增益
        if(infoGain>bestInfoGain):
            bestInfoGain=bestInfoGain
            bestFeature=i
    return bestFeature

# #测试
# myDat,lablels=createDataSet()
# chooseBestFeatureToSplit(myDat)
# print(myDat)