如何在自定义K-Fold交叉验证后将预测结果存入Pandas数据框新列

核心问题分析 #

你之前的赋值方法失效，本质是Pandas链式索引返回副本导致赋值不生效，加上需要逐次填充各测试折的结果到原数据框对应位置。以下是直接可行的解决方案：

步骤1：提前初始化目标列 #

在交叉验证开始前，先在原数据框中创建需要的新列，用NaN填充，方便后续逐折填充：

import numpy as np

# 根据需求创建列，示例包含预测结果、是否正确、置信度
df['prediction'] = np.nan
df['is_correct'] = np.nan
df['confidence'] = np.nan

步骤2：迭代中用.loc直接赋值 #

每次K折迭代时，拿到测试折的预测结果、真实标签及对应原数据框索引后，使用.loc直接定位原数据框的索引位置赋值，避免链式索引的副本问题：

# 假设每次迭代已获取以下变量：
# clf：训练好的模型
# X_test：测试集特征
# y_test：测试集真实标签
# idx：原数据框中对应测试折的索引列表（你已通过treatment逻辑获取）

# 1. 生成测试集预测结果
test_predictions = clf.predict(X_test)
# 2. 生成置信度（取预测概率的最大值）
test_confidences = clf.predict_proba(X_test).max(axis=1)

# 3. 赋值到原数据框对应位置
df.loc[idx, 'prediction'] = test_predictions
df.loc[idx, 'is_correct'] = (y_test.values == test_predictions)
df.loc[idx, 'confidence'] = test_confidences

关键注意事项 #

• 必须使用.loc[idx, 列名]赋值：df.iloc[idx]['new_col']这种链式索引会返回原数据的副本，赋值操作只修改副本，不会影响原数据框。
• 确保交叉验证折覆盖全量样本：如果最终数据框中还有NaN值，说明部分样本未被分配到任何测试折，需要检查你的自定义K-Fold逻辑。
• 避免索引重复赋值：如果你的treatment分组逻辑可能导致同一索引被多次分到测试折，需要根据业务需求处理（比如保留最后一次预测，或调整折的划分规则）。

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