id3算法实现决策树

import math
import copy
class ID3DTree(object):
    def __init__(self):
        self.tree={ }#生成的树
        self.dataSet = []#数据集
        self.labels={}#标签集
    #导入数据
    def loadDataSet(self,path,labels):
        recordlist = []
        fp = open(path,"rb")#读取文件内容
        content = fp.read().decode()
        fp.close()
        rowlist = content.splitlines()#按行转换成一维表
        recordlist = [row.split("\t") for row in rowlist if row.strip()]
        self.dataSet = recordlist
        self.labels = labels
    def train(self):
        labels = copy.deepcopy(self.labels)
        self.tree = self.buildTree(self.dataSet,labels)
    #构建决策树，创建决策树主程序
    def buildTree(self,dataSet,labels):
        #抽取源数据集的决策标签列
        cateList = [data[-1] for data in dataSet]
        #程序终止条件1：如果cateList只有一种决策标签，停止划分，返回这个决决策标签
        if cateList.count(cateList[0]) == len(cateList):
            return cateList[0]
        #程序终止条件2：如果数据集的第一个决策标签只有一个，则
        #返回这个决策标签
        if len(dataSet[0]) == 1:
            return self.maxCate(cateList)
        #算法核心:
        bestFeat = self.getBestFeat(dataSet) #返回数据集的最优特征值
        bestFeatLabel = labels[bestFeat]
        tree = {bestFeatLabel:{}}
        del(labels[bestFeat])
        #再次抽取最优特征轴的列向量
        uniqueVals = set([data[bestFeat] for data in dataSet])#去重
        for value in uniqueVals:#决策树的递归增长
            subLabels = labels[:] #将删除后的特征类别集建立子类别集
            #按最优特征列和值分隔数据集，即筛选出第bestFeat列值为value的所有元素，返回的是去掉这一列后剩下的数据
            #本例中根据bestFeat划分为买或者是不买这两个类别，作为左右子树，再继续划分
            splitDataset = self.splitDataSet(dataSet,bestFeat,value)
            subTree = self.buildTree(splitDataset,subLabels)
            tree[bestFeatLabel][value] = subTree
        return tree

    #计算出现次数最多的类别标签
    def maxCate(self,catelist):
        items = dict([(catelist.count(i),i) for  i in catelist])
        return items([max(items.keys())])

    #计算最优特征
    def getBestFeat(self,dataSet):
        #计算特征向量维，其中最后一列用于类别标签，因此要减去
        numFeatures = len(dataSet[0]) - 1#特征向量维数-1
        baseEntropy = self.computeEntropy(dataSet) #基础熵：源数据的香农熵,即买或者不买计算出来的熵值
        bestInfoGain = 0.0 #初始化最优的信息增益
        bestFeature = -1 #初始化最优的特征轴
        #外循环:遍历数据集各列，计算最优特征轴
        #i为数据集列索引：取值范围:0-(numFeatures-1)
        #依次遍历每个特征
        for i in range(numFeatures):
            uniqueVals = set([data[i] for data in dataSet])#去重：该列的唯一值集，也就是求类别，比如青年，中年，老年
            newEntropy = 0.0#初始化该列的香农熵
            #算不同类型中买与不买的比例以及熵
            for value in uniqueVals:#内循环按列和唯一值计算香农熵
                #按选定列i和唯一值分隔数据集，即筛选出第i列值为value的所有元素，返回的是去掉这一列后剩下的数据
                subDataSet = self.splitDataSet(dataSet,i,value)
                #各自在总体样本中所占的比例
                prob = len(subDataSet) / float(len(dataSet))
                #概率*对应小类别中买与不买的熵值
                newEntropy += prob * self.computeEntropy(subDataSet)
            infoGain = baseEntropy - newEntropy #计算最大增益
            if(infoGain > bestInfoGain):#如果信息增益>0
                bestInfoGain = infoGain #用当前信息增益代替之前的最优增益
                bestFeature = i #重置最优特征为当前列

        return bestFeature

    #计算信息熵
    def computeEntropy(self,dataSet):#计算香农熵
        datalen = float(len(dataSet))
        cateList = [data[-1] for data in dataSet]#从数据集中得到类别标签
        #得到类别为key,出现次数value的字典
        items = dict([(i,cateList.count(i)) for i in cateList])
        infoEntropy = 0.0 #初始化香农熵
        for key in items:#香农熵 = -plog2(p)
            prob = float(items[key])/datalen
            infoEntropy -= prob*math.log(prob,2)
        return infoEntropy

    #划分数据集，分隔数据集，删除特征轴所在的数据列，返回剩余的数据集
    def splitDataSet(self,dataSet,axis,value):
        rtnList = []
        for featVec in dataSet:
            if featVec[axis] == value:
                rFeatVec = featVec[:axis] #取0-（axis-1）的元素
                rFeatVec.extend(featVec[axis+1:])#将特征轴(列)之后的元素加回
                rtnList.append(rFeatVec)
        return rtnList

if __name__=='__main__':
    dtree=ID3DTree()
    dtree.loadDataSet("dataset.dat",["age","revenue","student","credit"])
    dtree.train()
    print(dtree.tree)

结果：