用决策树 对一个有10个标签2个类别的数据集 用python实现信息增益
时间: 2024-05-08 17:14:59 浏览: 150
以下是一个简单的用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,因此选择第二个属性作为决策树的根节点。
阅读全文