如何在训练集与测试集拆分后保留原始索引

问题本质 #

train_test_split本身会保留输入数据的索引信息，但如果你的X_train_res和y_train_res是numpy数组（无索引属性），拆分后就会自动生成从0开始的连续索引，丢失原始索引。

解决方法 #

1. 将numpy数组转回带索引的DataFrame/Series

如果你的预处理输出是numpy数组，先把它转换成DataFrame/Series并关联原始索引：

import pandas as pd

# 替换为你的原始特征列名和原始数据集索引
X_train_res_df = pd.DataFrame(X_train_res, columns=原特征列名列表, index=原数据集.index)
y_train_res_series = pd.Series(y_train_res, index=原数据集.index)

# 再执行拆分，此时会保留原始索引
X_train,X_test,y_train,y_test = train_test_split(X_train_res_df, y_train_res_series, test_size=0.25, random_state=42)

2. 直接基于原始DataFrame拆分（推荐）

如果还没做过采样等预处理，直接用带索引的原始DataFrame拆分，全程不转换成numpy数组：

# 假设df是原始数据集，target是目标列名称
X = df.drop('target', axis=1)
y = df['target']

# 拆分后的子集自动保留原始索引
X_train,X_test,y_train,y_test = train_test_split(X, y, test_size=0.25, random_state=42)

3. 过采样后立即关联索引（针对你的X_train_res场景）

如果X_train_res是过采样（比如SMOTE）后的结果，在过采样完成后立刻赋予原始索引：

from imblearn.over_sampling import SMOTE

sm = SMOTE(random_state=42)
# 基于带索引的X_train、y_train执行过采样
X_train_res, y_train_res = sm.fit_resample(X_train, y_train)

# 重新关联原始索引
X_train_res = pd.DataFrame(X_train_res, columns=X_train.columns, index=X_train.index)
y_train_res = pd.Series(y_train_res, index=X_train.index)

# 再拆分即可保留原始索引
X_train,X_test,y_train,y_test = train_test_split(X_train_res, y_train_res, test_size=0.25, random_state=42)

验证方式 #

拆分后可以通过print(X_train.index)查看是否保留了原始索引。

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