机器学习实战决策树（一）——信息增益与划分数据集

from math import log
#计算给定的熵
def calcsahnnonent(dataset):
        numentries = len(dataset)     #计算实例的总数
        labelcounts ={}
        #创建一个数据字典
        for featvec in dataset:
            currentlabel = featvec[-1]    #键值是最后一列数值                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                       
            if currentlabel not in labelcounts.keys():  #为所有可能的分类创建字典。使用的是字典中key（）方法
                labelcounts[currentlabel]= 0
            labelcounts[currentlabel] += 1
        shannonent = 0.0
        for key in labelcounts:
            prob = float(labelcounts[key]) / numentries
            shannonent -= prob * log(prob,2)    #以2为底求对数
        return shannonent
#创建数据集
def createdataset():
    dataset = [ [1, 1, 'yes'],
                [1, 1, 'yes'],
                [1, 0, 'no'],
                [0, 1, 'no'],
                [0, 1, 'no'] ]
    labels = ['no surfacing', 'flippers']
    return dataset, labels

#mydata,labels = createdataset()
#print('mydata:',mydata)
#mydata[0][-1] = 'maybe'
#calcsahnnonent(mydata)

#按照给定特征划分数据集
def splitdatset(dataset,axis,value):      #dataset:需要划分的数据集；axis:划分数据集的特征
    retdataset = []             #为了不修改原数据集，创建一个新列表对象
    for featvec in dataset:  #抽取数据
        if featvec[axis] == value:
            reducefeatvec = featvec[:axis]  #去掉axis 的特征
            #将符合条件的添加到返回的数据集
            reducefeatvec.extend(featvec[axis+1:]) 
            retdataset.append(reducefeatvec)
    return retdataset
#print(splitdatset(mydata,0,1))


#选择最好的数据集划分方式
def choosebestfeaturetosplit(dataset):
    numfeatures = len(dataset[0]) - 1
    baseentropy = calcsahnnonent(dataset)  #计算整个数据集的原始香农熵
    bestinfogain = 0.0   #初始化返回参数
    bestfeature = -1  #最优特征索引
    for i in range(numfeatures):   #遍历所有的特征
        featlist = [example[i] for example in dataset]   # 获取dataSet的第i个所有特征
        uniquevals = set(featlist)   #创建set集合{}，元素不可重复
        newentropy = 0.0 #经验条件熵初始化为0 
        #计算信息增益
        for value in uniquevals:  #遍历该特征的所有取值
            subdataset = splitdatset(dataset, i, value)
            prob = len(subdataset) / float(len(dataset))
            newentropy += prob * calcsahnnonent(subdataset)
        infogain = baseentropy - newentropy
        #寻找最大信息增益
        if (infogain > bestinfogain):
            bestinfogain = infogain
            bestfeature = i
    return bestfeature

choosebestfeaturetosplit(mydata)