机器学习-决策树 -计算信息增益

"""
文件说明：决策树
"""
from math import log

"""
函数说明：创建测试数据集
Parameters:
    null
Returns:
    dataSet - 数据集
    labels - 分类属性
Author：
    ZhengYuXiao
Modify：
    2019-03-04
"""


def createDataSet():
    dataSet = [[0, 0, 0, 0, 'no'],  # 数据集
               [0, 0, 0, 1, 'no'],
               [0, 1, 0, 1, 'yes'],
               [0, 1, 1, 0, 'yes'],
               [0, 0, 0, 0, 'no'],
               [1, 0, 0, 0, 'no'],
               [1, 0, 0, 1, 'no'],
               [1, 1, 1, 1, 'yes'],
               [1, 0, 1, 2, 'yes'],
               [1, 0, 1, 2, 'yes'],
               [2, 0, 1, 2, 'yes'],
               [2, 0, 1, 1, 'yes'],
               [2, 1, 0, 1, 'yes'],
               [2, 1, 0, 2, 'yes'],
               [2, 0, 0, 0, 'no']]
    labels = ['不放贷', '放贷']  # 分类属性
    return dataSet, labels  # 返回数据集和分类属性


"""
函数说明：计算给定数据集的经验熵
Paremeters：
    dataSet - 数据集
Returns：
    shannonEnt - 经验熵 H(D)
Author：
    ZhengYuXiao
Modofy：
    2019-03-04
"""


def calcShannonEnt(dataSet):
    # 训练集数据条数
    numEntires = len(dataSet)
    # 记录训练集中每个标签出现次数的字典
    labelCounts = {}
    for featVec in dataSet:
        currentLabel = featVec[-1]
        labelCounts[currentLabel] = labelCounts.get(currentLabel, 0) + 1
    # 经验熵
    shannonEnt = 0.0
    #         k
    # H(D) = -∑[P(Xi)log2P(Xi)]     P(Xi) = |Ck|/|D|
    #         i=1
    # |Ck|:训练集中标签为Ck的数据的条数
    # |D|:训练集数据总条数
    for key in labelCounts:
        prob = float(labelCounts[key]) / numEntires
        shannonEnt -= prob * log(prob, 2)
    return shannonEnt


"""
函数说明：按照给定特征以及指定的特征值划分数据集（例：性别 男）
Parameters：
    dataSet - 待划分的数据集
    axis - 划分数据集的特征的列数[0,1,2,3....]
    value - 需要返回的特征的值
Returns：
    returnDataSet - 指定特征的值等于指定特征值的数据的集合（例：性别为男的数据的集合）
Author：
    ZhengYuXiao
Modify：
    2019-03-05
"""


def splitDataSet(dataSet, axis, value):
    # 存储返回数据的列表
    returnDataSet = []
    # 遍历训练集
    for featVec in dataSet:
        # 如果该条数据中用来划分训练集的特征的值等于指定值
        # 则把这个特征值从这条数据中去除，并把新数据储存
        if featVec[axis] == value:
            # list[ start : end ] ：Python列表的切片操作
            # 返回列表从start位置开始，到end位置前的元素组成的列表
            #          例：print([1,3,5,7,9][1:3])
            #          输出：[3，5]
            #
            # list.extend( list2 ) :把列表list2中的元素添加到列表list的尾部
            #          例子：List1 = [0, 1, 2, 3, 4, 5, 6][:2]
            #                List1.extend([0, 1, 2, 3, 4, 5, 6][3:])
            #                print(List1)
            #          输出：[0, 1, 3, 4, 5, 6]
            reducedFeatVec = featVec[:axis]
            reducedFeatVec.extend(featVec[axis + 1:])
            returnDataSet.append(reducedFeatVec)
    return returnDataSet


"""
函数说明：选择划分训练集的最优特征
Parameters：
    dataSet - 数据集
Returns：
    bestFeature - 使信息增益最大的特征的索引值
Author：
    ZhengYuXiao
Modify：
    2019-03-05
"""


def chooseBestFeatureToSplit(dataSet):
    # 特征数量
    numFeatures = len(dataSet[0]) - 1
    # 训练集的香农熵 H(D)
    baseEntropy = calcShannonEnt(dataSet)
    # 信息增益
    bestInfoGain = 0.0
    # 最优特征（使训练集信息增益最大的特征）的索引值
    bestFeature = -1
    #
    # 信息增益 = 集合的熵 - 特征X给定条件下集合的熵
    # g(D,X) = H(D) - H(D|X)
    #
    # 计算每一个特征对集合产生的信息增益
    for i in range(numFeatures):
        # 把每一特征整列取出：
        #  性别  年级 及格     性别  年龄  及格
        # | 男    1    Y |     |男|  |1|  |Y|
        # | 女    4    Y |  -> |女|  |4|  |Y|
        # | 女    3    N |     |女|  |3|  |N|
        # | 男    1    Y |     |男|  |1|  |Y|
        featList = [example[i] for example in dataSet]
        # 创建集合，集合中的元素不能重复，set函数会自动去除重复的数据，
        # 所以集合存储——某一特征的所有取值
        uniqueVals = set(featList)
        # 条件熵 H(D|X)
        #                  n
        #               =  ∑[ |Di|/|D| * H(Di) ]
        #                  i=1
        #
        #                   n              K
        #               = - ∑[ |Di|/|D| * ∑[ |Dik|/|Di|] * Log2(|Dik|/|Di|) ]
        #                   i=1            k=1
        #
        newEntropy = 0.0
        for value in uniqueVals:
            # 训练集被某一特征的某一特征值划分的子集
            # 例：
            #     年级 及格
            #     | 1   Y |
            #     | 1   Y |
            subDataSet = splitDataSet(dataSet, i, value)
            # |Di|/|D|
            prob = len(subDataSet) / float(len(dataSet))
            # 根据公式计算条件经验熵 H(D|X)：
            #                  n
            #                  ∑[ |Di|/|D| * H(Di) ]
            #                  i=1
            newEntropy += prob * calcShannonEnt(subDataSet)
        # 计算信息增益：
        #     信息增益 = 集合的熵 - 特征X给定条件下集合的熵
        #     g(D,X) = H(D) - H(D|X)
        infoGain = baseEntropy - newEntropy
        print("第", i, "个特征的信息增益为：", infoGain)
        if (infoGain > bestInfoGain):
            bestInfoGain = infoGain
            bestFeature = i
    print("信息增益最高的是第", bestFeature, "个特征")
    return bestFeature


if __name__ == '__main__':
    dataSet, features = createDataSet()
    print("训练集=\n", dataSet)
    print("经验熵=\n", calcShannonEnt(dataSet))
    chooseBestFeatureToSplit(dataSet)