如何在scikit-learn中获取决策树的总加权基尼不纯度/熵？

问题背景 #

你需要计算训练好的决策树中所有叶子节点的加权不纯度之和（权重为叶子样本数占总样本数的比例），并对比决策树构建前后的总不纯度，但scikit-learn没有直接提供该指标的API。以下是基于泰坦尼克数据集的示例代码：

import pandas as pd
import matplotlib.pyplot as plt
from sklearn import tree
df_titanic = pd.read_csv('titanic_data.csv')    # 泰坦尼克数据集
feat_list = ['SibSp','Pclass']  # 特征：同乘的兄弟姐妹/配偶数量、客舱等级
clf = tree.DecisionTreeClassifier()
clf = clf.fit(df_titanic.loc[:,feat_list],df_titanic['Survived'])
fig = plt.figure(figsize=(10,10))
tree.plot_tree(clf,feature_names=feat_list,class_names=['NS','S'])
fig.show()

解决方案 #

scikit-learn的决策树模型可以通过内部属性tree_获取节点的关键信息，从而手动计算总加权不纯度：

1. 计算训练后的总加权基尼不纯度 #

import pandas as pd
from sklearn import tree

# 加载数据并训练决策树
df_titanic = pd.read_csv('titanic_data.csv')
feat_list = ['SibSp','Pclass']
clf = tree.DecisionTreeClassifier()
clf.fit(df_titanic.loc[:,feat_list], df_titanic['Survived'])

# 提取树的内部信息
tree_obj = clf.tree_
total_samples = df_titanic.shape[0]
weighted_gini = 0.0

# 遍历所有节点，筛选叶子节点计算加权和
for node_idx in range(tree_obj.node_count):
    # 叶子节点的判断：左右子节点索引均为-1
    if tree_obj.children_left[node_idx] == -1 and tree_obj.children_right[node_idx] == -1:
        node_gini = tree_obj.impurity[node_idx]
        node_sample_num = tree_obj.n_node_samples[node_idx]
        weighted_gini += node_gini * (node_sample_num / total_samples)

print(f"训练后的总加权基尼不纯度: {weighted_gini:.4f}")

2. 计算构建前的初始基尼不纯度 #

初始不纯度即根节点未分裂时的全局不纯度，可直接基于目标变量计算：

import numpy as np

def compute_gini(y):
    # 统计各类别样本数
    _, class_counts = np.unique(y, return_counts=True)
    # 计算各类别占比
    class_probs = class_counts / len(y)
    # 基尼不纯度公式：1 - Σ(p_i²)
    return 1 - np.sum(np.square(class_probs))

initial_gini = compute_gini(df_titanic['Survived'])
print(f"初始总基尼不纯度: {initial_gini:.4f}")

3. 针对熵准则的适配 #

如果训练时使用熵作为分裂准则（criterion='entropy'），只需调整计算逻辑：

# 训练基于熵的决策树
clf_entropy = tree.DecisionTreeClassifier(criterion='entropy')
clf_entropy.fit(df_titanic.loc[:,feat_list], df_titanic['Survived'])

# 计算总加权熵
tree_obj_entropy = clf_entropy.tree_
weighted_entropy = 0.0
for node_idx in range(tree_obj_entropy.node_count):
    if tree_obj_entropy.children_left[node_idx] == -1 and tree_obj_entropy.children_right[node_idx] == -1:
        node_entropy = tree_obj_entropy.impurity[node_idx]
        node_sample_num = tree_obj_entropy.n_node_samples[node_idx]
        weighted_entropy += node_entropy * (node_sample_num / total_samples)

# 计算初始熵
def compute_entropy(y):
    _, class_counts = np.unique(y, return_counts=True)
    class_probs = class_counts / len(y)
    # 过滤掉概率为0的项，避免log(0)报错
    class_probs = class_probs[class_probs > 0]
    # 熵公式：-Σ(p_i * log2(p_i))
    return -np.sum(class_probs * np.log2(class_probs))

initial_entropy = compute_entropy(df_titanic['Survived'])
print(f"初始总熵: {initial_entropy:.4f}")
print(f"训练后的总加权熵: {weighted_entropy:.4f}")

关键属性说明 #

• clf.tree_.impurity：存储每个节点的不纯度（基尼或熵，由训练时的criterion决定）
• clf.tree_.n_node_samples：存储每个节点包含的样本数量
• clf.tree_.children_left/children_right：每个节点的左/右子节点索引，叶子节点的子节点索引为-1

内容的提问来源于stack exchange，提问作者J Prestone