如何剪枝sklearn决策树以保证输出正样本概率单调递增

核心逻辑 #

由于输入为单特征，决策树的所有分裂都是按特征值从小到大划分区间，因此要求左子节点（特征≤阈值）的正样本概率必须≤右子节点（特征>阈值）的正样本概率，不符合该单调性要求的节点需要剪枝：

• 剪枝操作直接将冲突节点的父节点设置为叶子节点，采用父节点整体样本的正样本占比作为输出概率
• 若待取舍的叶子节点归属不同父节点，优先保留样本量更大的节点对应的子树结构

完整实现代码 #

from sklearn.tree import DecisionTreeClassifier, plot_tree
from sklearn.tree._tree import TREE_LEAF
import numpy as np
from sklearn.datasets import load_iris
import matplotlib.pyplot as plt

# 递归剪枝核心函数
def prune_non_monotonic(tree, node_id=0):
    left = tree.tree_.children_left[node_id]
    right = tree.tree_.children_right[node_id]
    
    # 叶子节点直接返回样本量、正概率、叶子标记
    if left == TREE_LEAF and right == TREE_LEAF:
        n_samples = tree.tree_.n_node_samples[node_id]
        pos_prob = tree.tree_.value[node_id, 0, 1] / np.sum(tree.tree_.value[node_id, 0])
        return n_samples, pos_prob, True
    
    # 自底向上递归处理左右子树
    left_n, left_prob, left_is_leaf = prune_non_monotonic(tree, left)
    right_n, right_prob, right_is_leaf = prune_non_monotonic(tree, right)
    
    # 左右子树均为叶子时检查单调性
    if left_is_leaf and right_is_leaf:
        # 符合单调要求直接返回
        if left_prob <= right_prob + 1e-8: # 增加浮点误差容限
            total_n = left_n + right_n
            avg_prob = (left_n * left_prob + right_n * right_prob) / total_n
            return total_n, avg_prob, False
        # 不符合要求则剪枝，将当前节点设为叶子
        else:
            tree.tree_.children_left[node_id] = TREE_LEAF
            tree.tree_.children_right[node_id] = TREE_LEAF
            n_samples = tree.tree_.n_node_samples[node_id]
            pos_prob = tree.tree_.value[node_id, 0, 1] / np.sum(tree.tree_.value[node_id, 0])
            return n_samples, pos_prob, True
    
    # 子树未完全剪为叶子时返回总样本和平均概率
    total_n = left_n + right_n
    avg_prob = (left_n * left_prob + right_n * right_prob) / total_n
    return total_n, avg_prob, False

# 基础训练逻辑
iris = load_iris()
X, y = iris.data[:, 0].reshape(-1,1), np.where(iris.target==0,0,1)
tree = DecisionTreeClassifier(max_depth=3, random_state=123)
tree.fit(X,y)

# 执行剪枝
prune_non_monotonic(tree)

# 验证单调性工具函数
def get_leaf_probs(tree, node_id=0, threshold_rules=None):
    threshold_rules = threshold_rules or []
    left = tree.tree_.children_left[node_id]
    right = tree.tree_.children_right[node_id]
    if left == TREE_LEAF and right == TREE_LEAF:
        prob = tree.tree_.value[node_id,0,1]/np.sum(tree.tree_.value[node_id,0])
        return [(" 且 ".join(threshold_rules), round(prob,4))]
    cur_thres = round(tree.tree_.threshold[node_id],2)
    left_probs = get_leaf_probs(tree, left, threshold_rules + [f"特征≤{cur_thres}"])
    right_probs = get_leaf_probs(tree, right, threshold_rules + [f"特征>{cur_thres}"])
    return left_probs + right_probs

# 输出剪枝结果
leaf_info = get_leaf_probs(tree)
print("剪枝后叶子节点正样本概率（按特征从小到大排列）：")
for rule, prob in leaf_info:
    print(f"{rule}: {prob}")

# 可视化剪枝后树结构
plt.figure(figsize=(12,6))
plot_tree(tree, filled=True, proportion=True, fontsize=10)
plt.show()

实现说明 #

• 递归采用自底向上的处理逻辑，优先修剪深层子树的非单调结构，再向上层传递结果，可适配任意深度的决策树
• 直接修改sklearn内置的树结构属性，无需重新训练模型，执行效率高
• 内置浮点误差容限，避免数值精度问题导致的误剪枝
• 剪枝后的所有叶子节点正样本概率严格随特征值增大单调递增，符合要求

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