决策树之信息增益计算

import numpy as np

def entropy(labels):
    """计算信息熵"""
    _, counts = np.unique(labels, return_counts=True)
    probabilities = counts / len(labels)
    return -np.sum(probabilities * np.log2(probabilities))

def conditional_entropy(feature, labels):
    """计算条件熵"""
    unique_values = np.unique(feature)
    total = len(labels)
    cond_ent = 0
    
    for value in unique_values:
        # 获取该特征值对应的所有标签
        subset = labels[feature == value]
        # 计算占比
        weight = len(subset) / total
        # 累加: 占比 * 该子集的熵
        cond_ent += weight * entropy(subset)
    
    return cond_ent

def information_gain(feature, labels):
    """计算信息增益"""
    # 信息熵
    ent = entropy(labels)
    # 条件熵
    cond_ent = conditional_entropy(feature, labels)
    # 信息增益 = 信息熵 - 条件熵
    return ent - cond_ent

# 1. 案例数据
# 6个样本，3个A，3个B
labels = np.array(['A', 'A', 'A', 'B', 'B', 'B'])
# 特征A：alpha(前4个), beta(后2个)
feature = np.array(['alpha', 'alpha', 'alpha', 'alpha', 'beta', 'beta'])

# 2. 计算信息熵
ent = entropy(labels)
print(f"信息熵 H(D): {ent:.4f}")  # 1.0

# 3. 计算条件熵
cond_ent = conditional_entropy(feature, labels)
print(f"条件熵 H(D|A): {cond_ent:.4f}")

# 4. 计算信息增益
ig = information_gain(feature, labels)
print(f"信息增益 g(D,A): {ig:.4f}")

# 5. 详细计算过程
print("\n=== 详细计算过程 ===")
print(f"数据: labels={labels}")
print(f"特征: feature={feature}")

# 信息熵
print(f"\n1. 信息熵 H(D):")
unique_labels, label_counts = np.unique(labels, return_counts=True)
for label, count in zip(unique_labels, label_counts):
    p = count / len(labels)
    print(f"  {label}: {count}/{len(labels)} = {p:.4f}, -{p}*log2({p}) = {-p * np.log2(p):.4f}")
print(f"  H(D) = {ent:.4f}")

# 条件熵
print(f"\n2. 条件熵 H(D|A):")
for value in np.unique(feature):
    subset = labels[feature == value]
    subset_ent = entropy(subset)
    weight = len(subset) / len(labels)
    print(f"  {value}: {len(subset)}/{len(labels)} = {weight:.4f}, H={subset_ent:.4f}")
    print(f"    贡献: {weight:.4f} * {subset_ent:.4f} = {weight * subset_ent:.4f}")
print(f"  H(D|A) = {cond_ent:.4f}")

# 信息增益
print(f"\n3. 信息增益:")
print(f"  g(D,A) = H(D) - H(D|A) = {ent:.4f} - {cond_ent:.4f} = {ig:.4f}")

# 6. 比较不同特征的信息增益
print("\n=== 比较不同特征的信息增益 ===")
features = {
    '特征1': np.array(['alpha', 'alpha', 'alpha', 'alpha', 'beta', 'beta']),
    '特征2': np.array(['alpha', 'alpha', 'beta', 'beta', 'alpha', 'alpha']),
    '特征3': np.array(['alpha', 'beta', 'alpha', 'beta', 'alpha', 'beta'])
}

for name, feat in features.items():
    ig = information_gain(feat, labels)
    print(f"{name}: 信息增益 = {ig:.4f}")