#-*-coding:utf-8-*-
from math import log
import operator
def calcShanonEnt(dataSet):
'''
計算給定數據集的香農熵
:param dataSet:
:return:shanonEnt
'''
numEntries = len(dataSet)
labelCounts={}
for featVec in dataSet:
currentLabel = featVec[-1]
if currentLabel not in labelCounts.keys():
labelCounts[currentLabel]=0
labelCounts[currentLabel] +=1
shanonEnt = 0.0
for key in labelCounts:
prob= float(labelCounts[key])/numEntries
shanonEnt -= prob*log(prob,2)
return shanonEnt
def splitDataSet(dataSet,axis,value):
'''
按照給特定特徵劃分數據集
:param dataSet:
:param axis:
:param value:
:return:
'''
retDataSet=[]
for featVec in dataSet:
if featVec[axis] == value:
reducedFeatVec = featVec[:axis]
reducedFeatVec.extend(featVec[axis+1:])
retDataSet.append(reducedFeatVec)
return retDataSet
def createDataSet():
'''
數據集
:return:
'''
dataSet = [
[1, 1, 'yes'],
[1, 1, 'yes'],
[1, 0, 'no'],
[0, 1, 'no'],
[0, 1, 'no'],
]
labels =['no surfacing','filppers']
return dataSet,labels
def chooseBestFeatureToSplit(dataSet):
'''
選擇最好的特徵進行分割
:param dataSet:
:return:
'''
numFeatures = len(dataSet[0])-1
baseEntropy = calcShanonEnt(dataSet)
bestInfoGain=0.0;bestFeature=-1
for i in range(numFeatures):
featList = [example[i] for example in dataSet]
uniqueVals =set(featList)
newEntropy = 0.0
for value in uniqueVals:
subDataSet = splitDataSet(dataSet,i,value)
prob = len(subDataSet)/float(len(dataSet))
newEntropy+= prob*calcShanonEnt(subDataSet)
infoGain =baseEntropy-newEntropy
if infoGain>bestInfoGain:
bestInfoGain =infoGain
bestFeature=i
return bestFeature
def majortyCnt(classList):
classCount={}
for vote in classList:
if vote not in classCount.keys():classCount[vote]=0
classCount[vote]+=1
sortedClassCount = sorted(classCount.iteritems(),key=operator.itemgetter(1),reverse=True)
return sortedClassCount[0][0]
def createTree(dataSet,labels):
'''
創建樹的函數代碼
:param dataSet:
:param labels:
:return:
'''
classList = [example[-1] for example in dataSet]
if classList.count(classList[0]) == len(classList):
return classList[0]
if len(dataSet[0]) == 1:
return majortyCnt(classList)
bestFeat = chooseBestFeatureToSplit(dataSet)
bestFeatLabel = labels[bestFeat]
myTree = {bestFeatLabel:{}}
del(labels[bestFeat])
featValues = [example[bestFeat] for example in dataSet]
uniqueVals = set(featValues)
for value in uniqueVals:
subLabels =labels[:]
myTree[bestFeatLabel][value]=createTree(splitDataSet(dataSet,bestFeat,value),subLabels)
return myTree
if __name__ == '__main__':
myDat,labels =createDataSet()
#print(myDat)
#print(calcShanonEnt(myDat))
#print(splitDataSet(myDat,0,1))
#print(splitDataSet(myDat, 0, 0))
#print(chooseBestFeatureToSplit(myDat))
#print(myDat)
print(createTree(myDat,labels))
作者:WangB