李航统计学习方法决策树的特征选择例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)