PySpark决策树回归模型：如何获取单条样本的决策路径？

PySpark ML库本身没有直接提供获取单条样本决策路径的API，但可以通过解析模型的树结构来实现这个功能。下面结合你的代码一步步说明：

1. 先修正代码中的笔误 #

你的代码最后一行clf_best = clf.bestModel里的clf是未定义的，应该改成cv_model.bestModel：

clf_best = cv_model.bestModel  # 这才是训练好的最优决策树回归模型

2. 解析决策树结构 #

训练好的DecisionTreeRegressionModel对象有两个关键方法可以获取树结构：

• toDebugString()：输出人类可读的树结构文本
• json()：返回结构化的JSON格式树数据，更适合程序解析

我们可以基于json()返回的结构来编写函数，遍历树并跟踪单条样本的决策路径。

3. 实现获取决策路径的函数 #

下面是一个示例函数，输入模型和单条样本（Row对象），返回该样本的决策路径：

def get_decision_path(tree_model, sample):
    # 从模型的JSON结构中提取树的根节点
    tree_json = tree_model.json()
    root_node = tree_json['paramMap']['ml_tree']['rootNode']
    
    path = []
    current_node = root_node
    
    while 'leftChild' in current_node:
        # 获取当前节点的分裂特征和分裂阈值
        split_col = current_node['split']['featureIndex']
        split_threshold = current_node['split']['threshold']
        
        # 获取样本对应特征的值
        feature_value = sample[split_col]
        
        # 判断走左分支还是右分支
        if feature_value <= split_threshold:
            direction = "left"
            path.append(f"Feature {split_col} <= {split_threshold} → left")
            current_node = current_node['leftChild']
        else:
            direction = "right"
            path.append(f"Feature {split_col} > {split_threshold} → right")
            current_node = current_node['rightChild']
    
    # 到达叶子节点，记录预测值
    path.append(f"Leaf node, predicted value: {current_node['prediction']}")
    
    return path

4. 使用示例 #

假设test_sample是你要查看的单条样本（从测试集中取一条，比如test.first()），调用函数即可得到决策路径：

test_sample = test.first()
decision_path = get_decision_path(clf_best, test_sample)

# 打印路径
for step in decision_path:
    print(step)

注意事项 #

• 函数中使用特征索引来获取样本值，如果你需要特征名称，可以结合模型的featureCol对应的VectorAssembler（如果有的话）来映射索引到名称
• 对于分类树，逻辑类似，只是分裂判断和叶子节点的含义不同

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