李航统计学习方法决策树的特征选择例5.2python代码

自己写一直报错一直报错。。。然后还是参考了机器学习实战上关于数据集划分的部分，所以这篇主要为了记录一下实战上的好方法：还可以这样分？！

先上图：

下面直接上代码：

# -*- coding:utf-8 -*-
#@author:xinxinzhang
from math import log
def loadDataSet():
    dataSet = [['青年', '否', '否', '一般', '否'],
               ['青年', '否', '否', '好', '否'],
               ['青年', '是', '否', '好', '是'],
               ['青年', '是', '是', '一般', '是'],
               ['青年', '否', '否', '一般', '否'],
               ['中年', '否', '否', '一般', '否'],
               ['中年', '否', '否', '好', '否'],
               ['中年', '是', '是', '好', '是'],
               ['中年', '否', '是', '非常好', '是'],
               ['中年', '否', '是', '非常好', '是'],
               ['老年', '否', '是', '非常好', '是'],
               ['老年', '否', '是', '好', '是'],
               ['老年', '是', '否', '好', '是'],
               ['老年', '是', '否', '非常好', '是'],
               ['老年', '否', '否', '一般', '否']]
    label = ['年龄', '有工作', '有自己的房子', '信贷情况']
    return dataSet, label

def shannoEnt(dataSet):
    labelCount={}
    numOfData=len(dataSet)
    for data in dataSet:
        classify=data[-1]
        if classify not in labelCount.keys():
            labelCount[classify]=1
        else:labelCount[classify]+=1
    H=0.0
    for value in labelCount.values():
        pi=value/numOfData
        H-=pi*log(pi,2)
    return H

def chooseBestFeatureSplit(dataSet):
    HD=shannoEnt(dataSet)
    bestGain=0.0
    bestFeature=-1   #为啥初始化为1，因为特征从0开始扫描，所以如果没找到最优特征，就会返回-1
    for i in range(len(dataSet[0])-1):
        feat=[data[i] for data in dataSet]
        prob=0.0
        for value in set(feat):
            subData=splitData(dataSet,i,value)
            prob+=(len(subData)/len(dataSet))*shannoEnt(subData)
        Gain=HD-prob
        if Gain>bestGain:
            bestGain=Gain
            bestFeature=i
    print('最优特征是：',label[bestFeature])
    return bestFeature


def splitData(dataSet,axis,value):  #机器学习实战上关于数据集的划分
    retDataSet=[]
    for data in dataSet:
        if data[axis]==value:
            reducedData=data[:axis]
            reducedData.extend(data[axis+1:])
            retDataSet.append(reducedData)
    return retDataSet



dataSet,label=loadDataSet()
chooseBestFeatureSplit(dataSet)