如何从CatBoost模型获取类似XGBoost的树结构DataFrame？

CatBoost没有提供和XGBoost trees_to_dataframe()完全等效的API，但可以通过解析模型的树结构信息，手动整理成符合需求的DataFrame。以下是具体实现方案：

实现步骤与代码示例 #

1. 导入依赖库并训练模型 #

import pandas as pd
from catboost import CatBoostRegressor
from catboost.utils import get_tree_structure

# 假设X['train']、Y['train']、n_num_features、n_features已提前定义
model = CatBoostRegressor(iterations=10, thread_count=10)
model.fit(X['train'], Y['train'], cat_features=range(n_num_features, n_features))

2. 编写转换函数 #

通过get_tree_structure()获取每棵树的节点信息，递归遍历后整理成目标DataFrame：

def catboost_trees_to_dataframe(model):
    trees_data = []
    # 遍历每一棵决策树
    for tree_idx in range(model.tree_count_):
        tree = get_tree_structure(model, tree_idx)
        
        # 递归遍历树的所有节点
        def traverse_node(node, node_id):
            # 构造当前节点的DataFrame行
            row = {
                'Tree': tree_idx,
                'Node': node_id,
                'ID': f"{tree_idx}-{node_id}",
                'Feature': node.split_feature if node.split_feature is not None else None,
                'Split': node.split_value if node.split_value is not None else None,
                'Yes': node.left_child if node.left_child is not None else None,
                'No': node.right_child if node.right_child is not None else None,
                'Missing': node.default_child if node.default_child is not None else None,
                'Gain': node.split_gain if node.split_gain is not None else None,
                'Cover': node.data.get('weight')  # 对应样本权重和，与XGBoost的Cover含义一致
            }
            trees_data.append(row)
            
            # 递归处理左、右子节点
            if node.left_child is not None:
                traverse_node(tree[node.left_child], node.left_child)
            if node.right_child is not None:
                traverse_node(tree[node.right_child], node.right_child)
        
        # 从根节点（ID=0）开始遍历
        traverse_node(tree[0], 0)
    
    return pd.DataFrame(trees_data)

3. 生成目标DataFrame #

trees_df = catboost_trees_to_dataframe(model)

列字段说明 #

生成的DataFrame包含你要求的所有字段：

• Tree: 决策树的索引（从0开始）
• Node: 当前节点在所属树中的ID（根节点为0）
• ID: 树索引+节点ID的唯一标识
• Feature: 节点分裂使用的特征（叶节点为None）
• Split: 分裂阈值（类别特征显示内部编码值，叶节点为None）
• Yes: 满足分裂条件时跳转的子节点ID
• No: 不满足分裂条件时跳转的子节点ID
• Missing: 缺失值对应的默认子节点ID
• Gain: 分裂带来的增益值（叶节点为None）
• Cover: 节点覆盖的样本权重和，与XGBoost的Cover字段含义一致

补充说明 #

• 若需要将类别特征的编码值转换为原始类别名称，可以结合model.feature_names_和model.get_cat_feature_info()进行映射。
• get_tree_structure()返回的树结构包含节点的详细元数据，可根据需求扩展更多字段。

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