AI4 决策树的生成与训练-信息增益

# -*- coding: UTF-8 -*-
from math import log
import pandas as pd

dataSet = pd.read_csv('dataSet.csv', header=None).values.tolist()


# 给定一个数据集，calcInfoEnt可以用于计算一个数据集的信息熵，可直接调用
# 也可不使用，通过自己的方式计算信息增益
def calcInfoEnt(data):
    numEntres = len(data)
    labelcnt = {}  # 用于统计正负样本的个数
    for item in data:
        if item[-1] not in labelcnt:
            labelcnt[item[-1]] = 0
        labelcnt[item[-1]] += 1
    infoEnt = 0.0
    for item in labelcnt:  # 根据信息熵的公式计算信息熵
        curr_info_entr = float(labelcnt[item]) / numEntres
        infoEnt = infoEnt - curr_info_entr * log(curr_info_entr, 2)
    return infoEnt
    # 返回值 infoEnt 为数据集的信息熵


# 给定一个数据集，用于切分一个子集，可直接用于计算某一特征的信息增益
# 也可不使用，通过自己的方式计算信息增益
# dataSet是要划分的数据集，i 代表第i个特征的索引index
# value对应该特征的某一取值
def create_sub_dataset(dataSet, i, value):
    res = []
    for item in dataSet:
        if item[i] == value:
            curr_data = item[:i] + item[i + 1:]
            res.append(curr_data)
    return res


def calc_max_info_gain(dataSet):  # 计算所有特征的最大信息增益，dataSet为给定的数据集
    n = len(dataSet[0]) - 1  # n 是特征的数量，-1 的原因是最后一列是分类标签
    total_entropy = calcInfoEnt(dataSet)  # 整体数据集的信息熵
    max_info_gain = [0, 0]  # 返回值初始化
    # code start here
    best_feature=-1
    for i in range(n):
        featList=[example[i] for example in dataSet]
        uniqueVals=set(featList)
        newEntropy=0.0
        for value in uniqueVals:
            subDataSet=create_sub_dataset(dataSet,i,value)
            prob=len(subDataSet)/float(len(dataSet))
            newEntropy+=prob*calcInfoEnt(subDataSet)
        infoGain=total_entropy-newEntropy
        if(infoGain>max_info_gain[1]):
            max_info_gain[1]=infoGain
            max_info_gain[0]=i
            best_feature=1

    # code end here
    return max_info_gain


if __name__ == '__main__':
    info_res = calc_max_info_gain(dataSet)
    print("信息增益最大的特征索引为：{0},对应的信息增益为{1}".format(info_res[0], info_res[1]))