Python实现决策树算法ID3

#coding=utf-8

# 构造数据集
def createDataSet():  
    dataSet = [[1,1,'yes'],[1,1,'yes'],[1,0,'no'],[0,1,'no'],[0,1,'no']]  
    features = ['no surfacing','flippers']  
    return dataSet,features  


# 采用字典来递归构造决策树
def treeGrowth(dataSet,features):  
    classList = [example[-1] for example in dataSet]  
    if classList.count(classList[0])==len(classList):  
        return classList[0]  
    if len(dataSet[0])==1:# no more features  
        return classify(classList)  
  
    bestFeat = findBestSplit(dataSet)#bestFeat is the index of best feature  
    bestFeatLabel = features[bestFeat]  
    myTree = {bestFeatLabel:{}}  
    featValues = [example[bestFeat] for example in dataSet]  
    uniqueFeatValues = set(featValues)  
    del (features[bestFeat])  
    for values in uniqueFeatValues:  
        subDataSet = splitDataSet(dataSet,bestFeat,values)  
        myTree[bestFeatLabel][values] = treeGrowth(subDataSet,features)  
    return myTree  


# 当没有多余的feature，但是剩下的样本不完全是一样的类别是，采用出现次数多的那个类别
def classify(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 findBestSplit(dataset):  
    numFeatures = len(dataset[0])-1  
    baseEntropy = calcShannonEnt(dataset)  
    bestInfoGain = 0.0  
    bestFeat = -1  
    for i in range(numFeatures):  
        featValues = [example[i] for example in dataset]  
        uniqueFeatValues = set(featValues)  
        newEntropy = 0.0  
        for val in uniqueFeatValues:  
            subDataSet = splitDataSet(dataset,i,val)  
            prob = len(subDataSet)/float(len(dataset))  
            newEntropy += prob*calcShannonEnt(subDataSet)  
        if(baseEntropy - newEntropy)>bestInfoGain:  
            bestInfoGain = baseEntropy - newEntropy  
            bestFeat = i  
    return bestFeat  


# 选择完分裂属性后，就行数据集的分裂
def splitDataSet(dataset,feat,values):  
    retDataSet = []  
    for featVec in dataset:  
        if featVec[feat] == values:  
            reducedFeatVec = featVec[:feat]  
            reducedFeatVec.extend(featVec[feat+1:])  
            retDataSet.append(reducedFeatVec)  
    return retDataSet


# 计算数据集的熵
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  
        if prob != 0:  
            shannonEnt -= prob*log(prob,2)  
    return shannonEnt  


# 根据构造的决策树进行数据的分类
def predict(tree,newObject):  
    while isinstance(tree,dict):  
        key = tree.keys()[0]  
        tree = tree[key][newObject[key]]  
    return tree  
  
if __name__ == '__main__':  
    dataset,features = createDataSet()  
    tree = treeGrowth(dataset,features)  
    print tree  
    print predict(tree,{'no surfacing':1,'flippers':1})  
    print predict(tree,{'no surfacing':1,'flippers':0})  
    print predict(tree,{'no surfacing':0,'flippers':1})  
    print predict(tree,{'no surfacing':0,'flippers':0})