题解 | #决策树的生成与训练-信息熵的计算#

import pandas as pd 
import numpy as np 

# 计算节点的信息熵：传入标签集(np.array)
def calcInfoEnt(y):
    H = 0
    if len(y) != 0:
        # 计算概率(纯度)
        p = len(y[y == 0]) / len(y) 
        # 计算信息熵(二分类公式)
        if p != 0 and p != 1:       # 如果概率为1，直接返回信息熵为0(无混乱程度无需分裂)
            H = -p * np.log2(p) - (1-p) * np.log2(1-p)          
    
    return H

# 取出数据集
dataSet = pd.read_csv('dataSet.csv',header=None)

# 分割特征集(x)和标签集(y)
x , y = dataSet.iloc[:,:-1].values , dataSet.iloc[:,-1].values 

# 清洗标签集
y = np.array([int(x.strip("'").strip(" '")) for x in y])

# 调用函数计算信息熵
infoEnt = calcInfoEnt(y)
print(infoEnt)