CART决策树代码实现：从算法原理到工程实践

1次阅读

共计 3721 个字符，预计需要花费 10 分钟才能阅读完成。

决策树是机器学习中最直观的算法之一，它通过树形结构对数据进行分类或回归。在众多决策树算法中，ID3、C4.5 和 CART 是最经典的三种。它们的主要区别在于特征选择标准和处理数据类型的能力：

ID3：使用信息增益选择特征，只能处理离散特征，容易过拟合
C4.5：改进 ID3，使用信息增益比，能处理连续特征，引入剪枝
CART：使用基尼系数，可同时处理分类和回归任务，二叉树结构更高效

基尼系数衡量数据的不纯度，计算公式为：

Gini(D) = 1 - Σ(p_i)^2

其中 p_i 是第 i 类样本在数据集 D 中的比例。对于特征 A 将数据集分为 D1 和 D2 后的基尼指数：

Gini_index(A) = |D1|/|D|*Gini(D1) + |D2|/|D|*Gini(D2)

递归构建决策树的终止条件包括：

当前节点所有样本属于同一类别
没有更多特征可用于分割
达到预设的最大深度
节点样本数小于预设阈值

后剪枝策略包括：

代价复杂度剪枝：通过平衡误差和复杂度选择最优子树
悲观剪枝：使用统计检验判断节点是否需要合并
MDL 原则：基于最小描述长度原则进行剪枝

from typing import List, Dict, Union, Optional
import numpy as np

class TreeNode:
    """决策树节点类"""
    def __init__(self, feature: Optional[int] = None, threshold: Optional[float] = None,
                 value: Optional[Union[int, float]] = None, left: Optional['TreeNode'] = None,
                 right: Optional['TreeNode'] = None):
        self.feature = feature  # 分裂特征索引
        self.threshold = threshold  # 分裂阈值
        self.value = value  # 叶节点预测值
        self.left = left  # 左子树
        self.right = right  # 右子树

class CARTDecisionTree:
    """CART 决策树实现"""
    def __init__(self, max_depth: int = 5, min_samples_split: int = 2):
        self.max_depth = max_depth
        self.min_samples_split = min_samples_split
        self.root: Optional[TreeNode] = None

    def _gini(self, y: np.ndarray) -> float:
        """计算基尼系数"""
        _, counts = np.unique(y, return_counts=True)
        probabilities = counts / counts.sum()
        return 1 - np.sum(probabilities ** 2)

    def _best_split(self, X: np.ndarray, y: np.ndarray) -> tuple:
        """寻找最佳分割特征和阈值"""
        best_gini = float('inf')
        best_feature, best_threshold = None, None

        for feature in range(X.shape[1]):
            thresholds = np.unique(X[:, feature])
            for threshold in thresholds:
                left_indices = X[:, feature] <= threshold
                gini = (left_indices.sum() * self._gini(y[left_indices]) + 
                       (~left_indices).sum() * self._gini(y[~left_indices])) / len(y)

                if gini < best_gini:
                    best_gini = gini
                    best_feature = feature
                    best_threshold = threshold

        return best_feature, best_threshold

    def _build_tree(self, X: np.ndarray, y: np.ndarray, depth: int = 0) -> TreeNode:
        """递归构建决策树"""
        # 终止条件检查
        if (depth >= self.max_depth or 
            len(y) < self.min_samples_split or 
            len(np.unique(y)) == 1):
            return TreeNode(value=np.bincount(y).argmax())

        # 寻找最佳分割
        feature, threshold = self._best_split(X, y)
        if feature is None:
            return TreeNode(value=np.bincount(y).argmax())

        # 递归构建子树
        left_indices = X[:, feature] <= threshold
        left = self._build_tree(X[left_indices], y[left_indices], depth + 1)
        right = self._build_tree(X[~left_indices], y[~left_indices], depth + 1)

        return TreeNode(feature=feature, threshold=threshold, left=left, right=right)

    def fit(self, X: np.ndarray, y: np.ndarray) -> None:
        """训练决策树"""
        self.root = self._build_tree(X, y)

    def predict(self, X: np.ndarray) -> np.ndarray:
        """预测"""
        return np.array([self._predict_one(x, self.root) for x in X])

    def _predict_one(self, x: np.ndarray, node: TreeNode) -> Union[int, float]:
        """单样本预测"""
        if node.value is not None:
            return node.value
        if x[node.feature] <= node.threshold:
            return self._predict_one(x, node.left)
        return self._predict_one(x, node.right)

我们使用 UCI 的鸢尾花数据集进行测试：

from sklearn.datasets import load_iris
from sklearn.tree import DecisionTreeClassifier
from sklearn.model_selection import train_test_split

# 加载数据
iris = load_iris()
X_train, X_test, y_train, y_test = train_test_split(iris.data, iris.target, test_size=0.2)

# 自定义模型
our_tree = CARTDecisionTree(max_depth=3)
our_tree.fit(X_train, y_train)
our_score = (our_tree.predict(X_test) == y_test).mean()

# sklearn 模型
sk_tree = DecisionTreeClassifier(criterion='gini', max_depth=3)
sk_tree.fit(X_train, y_train)
sk_score = sk_tree.score(X_test, y_test)

print(f"Our CART accuracy: {our_score:.2f}")
print(f"sklearn accuracy: {sk_score:.2f}")

通过绘制不同 max_depth 下的训练 / 测试准确率曲线，可以观察到：