决策树ID3的实现-python

import numpy as np
import pandas as pd

# 创造数据集
def createDataSet():
    data_row = {
        'no surfacing': [1, 1, 1, 0, 0],
        'flippers': [1, 1, 0, 1, 1],
        'fish': ['yes', 'yes', 'no', 'no', 'no']
    }
    dataSet = pd.DataFrame(data_row)
    return dataSet

# 计算香农熵，为了计算熵
def calcShannonEnt(dataSet):
    '''
    计算香农熵
    :param dataSet: 原始数据集
    :return: ent:香农熵的值
    '''
    n = dataSet.shape[0]  # 数据集总行数
    iset = dataSet.iloc[:, -1].value_counts()  # 标签的所有类别
    p = iset / n  # 每一类标签所占比
    ent = (-p * np.log2(p)).sum()  # 计算信息熵
    return ent

# 按照给定的特征划分数据集
def splitDataSet(dataSet, axis, value):
    '''
    :param dataSet: 待划分数据集
    :param axis: 指定的列序号
    :param value: 指定的属性的值
    :return: 按照指定的列索引和属性值切分后的数据集
    '''
    col = dataSet.columns[axis]  # 获得该序号对应的特征名
    rrdataSet = dataSet.loc[dataSet[col] == value, :]  # 获得属性col==value时的数据集
    redataSet = rrdataSet.drop(col, axis=1)  # 去掉指定的划分属性那一列
    return redataSet

# 选择最优的列进行切分
def chooseBestFeatureToSplit(dataSet):
    baseShannon = calcShannonEnt(dataSet)  # 计算初始的香农熵
    bestGain = 0  # 初始化信息增益
    bestFeatureIndex = -1  # 初始化最佳切分列
    for i in range(dataSet.shape[1] - 1):#对特征的每一列进行循环
        values=dataSet.iloc[:,i].value_counts().index#取出当前列的所有取值
        newShannon = 0.0
        for value in values:
            subDataSet = splitDataSet(dataSet, i, value)  # 利用不同的特征以及该特征的不同的取值划分数据集
            # 求新划分的数据集的香农熵
            prob = float(len(subDataSet)) / float(len(dataSet))
            newShannon += prob * calcShannonEnt(subDataSet)
        infoGain = baseShannon - newShannon  # 计算当前列的信息增益
        print(dataSet.columns[i],'划分数据集后的信息增益为:',infoGain)
        if (infoGain > bestGain):
            bestGain = infoGain
            bestFeatureIndex = i
    print('最高的信息增益为：',bestGain)
    return bestFeatureIndex
'''
# 没用到！
# 多数表决法定义叶子结点的分类
import operator
def majorityCnt(classList):
    classCount = {}
    for vote in classList:
        if vote not in classCount.keys():
            classCount[vote] = 1
        else:
            classCount[vote] += 1
    sortedClassCount = sorted(classCount.items(), key=operator.itemgetter(1), reverse=True)
    return sortedClassCount
'''


# 递归构建决策树
def createTree(dataSet, featuresNames):
    print('========开始计算该结点是否需要划分========')
    print('目前的数据集为：')
    print(dataSet)
    print('目前的特征为：:',featuresNames)
    classList = dataSet.iloc[:, -1].value_counts()#最后一列的类标签的统计
    #print(classList.index)
    #print(classList.index[0])
    decisionTree={}
    #递归出口：判断最多的标签数目(no 是3)是否等于数据集行数，或者数据集是否只有一列
    if classList[0] == dataSet.shape[0] or dataSet.shape[1] == 1:
        print('==========此处不需要划分！==========')
        return classList.index[0]  # 返回该类标签
    #开始创建决策树
    bestFeatureIndex=chooseBestFeatureToSplit(dataSet)#选择出划分数据集最好的特征的序号
    bestFeatureName=featuresNames[bestFeatureIndex]#根据序号获得属性名
    print('划分数据集最好的特征为：',  bestFeatureName)
    decisionTree={bestFeatureName:{}}#将该特征作为树的根节点
    del featuresNames[bestFeatureIndex]#将已放进树中的特征从特征中删除
    bestFeatureValues=set(dataSet.iloc[:, bestFeatureIndex])#提取所有样本关于这个特征的取值的不同取值
    #print('删除已放进树中的特征后，特征集合为：',featuresNames)
    #print('所有样本关于这个特征的取值:',bestFeatureValues)
    #递归式
    for value in bestFeatureValues:
        subFeaturesNames=featuresNames[:]#前几行中已经删除了最佳特征
        decisionTree[bestFeatureName][value]=\
            createTree(splitDataSet(dataSet,bestFeatureIndex,value),subFeaturesNames)
    return decisionTree


#使用决策树对一个实例分类
def classify(inputTree, attributes, testVector):
    '''
    :param inputTree: 已经生成的决策树
    :param attributes:训练集的所有属性（所有列）
    :param testVector: 测试数据，顺序对应训练集
    :return: 分类结果
    '''
    #firstNode,=inputTree.keys()#获取决策树第一个节点
    firstNode=next(iter(inputTree))
    #print('决策树第一个节点：',firstNode)
    secondDict=inputTree[firstNode]#下一个字典
    #print('secondDict：', secondDict)
    featureIndex=attributes.index(firstNode)#第一个结点所在的索引
    for key in secondDict.keys():
        if testVector[featureIndex]==key:
            if type(secondDict[key]).__name__=='dict':
                classLabel=classify(secondDict[key],attributes,testVector)
            else:
                classLabel=secondDict[key]
    return classLabel


if __name__ == '__main__':
    dataSea=createDataSet()#数据集
    attributes = list(dataSea.columns) #所有属性
    featuresNames=attributes[0:2]#特征
    train = dataSea#训练集
    decisionTree=createTree(dataSea,featuresNames)#建树
    np.save('models/decisionTree.npy',decisionTree)#保存树模型
    #loadDecisionTree=np.load('models/decisionTree.npy',allow_pickle=True).item()#保存之后可以直接加载模型使用
    print('========测试=======')
    # 对一个实例分类
    res1 = classify(decisionTree, attributes, [0,0])
    # 对一个数据集分类
    testVectors=[[0,0],[0,1],[1,0],[1,1]]
    result = []
    for test in testVectors:
        newres = classify(decisionTree, list(train.columns), test)
        result.append(newres)
    print(result)