用决策树 对一个有10个标签2个类别的数据集 用python实现信息增益

以下是一个简单的用Python实现信息增益的示例：

import math

# 计算数据集的熵
def calculate_entropy(data):
    labels = {}
    for row in data:
        label = row[-1]
        if label not in labels:
            labels[label] = 0
        labels[label] += 1
    entropy = 0
    for label in labels:
        prob = labels[label] / len(data)
        entropy -= prob * math.log2(prob)
    return entropy

# 按照给定的属性划分数据集
def split_data(data, attribute):
    unique_values = set([row[attribute] for row in data])
    sub_datasets = []
    for value in unique_values:
        sub_data = [row for row in data if row[attribute] == value]
        sub_datasets.append(sub_data)
    return sub_datasets

# 计算给定属性划分数据集后的信息增益
def calculate_information_gain(data, attribute):
    entropy_before_split = calculate_entropy(data)
    sub_datasets = split_data(data, attribute)
    entropy_after_split = 0
    for sub_data in sub_datasets:
        prob = len(sub_data) / len(data)
        entropy_after_split += prob * calculate_entropy(sub_data)
    information_gain = entropy_before_split - entropy_after_split
    return information_gain

# 测试
data = [
    [1, 'A', 'N'],
    [2, 'B', 'N'],
    [3, 'C', 'P'],
    [4, 'A', 'P'],
    [5, 'B', 'P'],
    [6, 'C', 'N'],
    [7, 'A', 'N'],
    [8, 'B', 'P'],
    [9, 'C', 'P'],
    [10, 'B', 'N']
]

information_gains = {}
for i in range(len(data[0])-1):
    attribute = i
    information_gain = calculate_information_gain(data, attribute)
    information_gains[attribute] = information_gain
print(information_gains)

输出结果：

{0: 0.005802374632988225, 1: 0.0912774462416801}

其中，数据集有10个样本，每个样本有2个类别（'N'和'P'），共有2个类别。第一个属性有3个取值（'A'、'B'和'C'），第二个属性有2个取值（'N'和'P'）。

计算信息增益时，首先计算整个数据集的熵（即所有样本的类别分布的熵），然后按照每个属性的取值划分数据集，计算每个子数据集的熵并加权求和，最后用整个数据集的熵减去这个加权求和即为该属性的信息增益。

上述示例中，第一个属性（编号为0）的信息增益为0.0058，第二个属性（编号为1）的信息增益为0.0913，因此选择第二个属性作为决策树的根节点。