使用train_test_split拆分数据集后丢失列名与索引名该如何恢复

你遇到的问题根源是.values属性会将pandas的DataFrame/Series转换为不存储列名、索引等元信息的numpy数组，拆分后自然丢失了原有结构信息，可按以下两种场景处理：

场景1：尚未执行拆分，优先使用该方案 #

train_test_split原生支持直接传入pandas数据结构，不需要转成numpy数组，拆分后会完整保留所有元信息：

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size = 0.20, random_state=99)

执行后可直接通过X_train.columns获取原列名，X_train.index、y_train.index获取对应原数据集的行索引。

场景2：已使用.values完成拆分，手动恢复元信息 #

恢复列名 #

直接复用原X的columns属性，将numpy数组转回DataFrame时指定列名即可：

import pandas as pd
# 为训练集、测试集恢复列名
X_train = pd.DataFrame(X_train, columns=X.columns)
X_test = pd.DataFrame(X_test, columns=X.columns)

恢复行索引 #

如果需要保留和原数据集对齐的行索引，建议先拆分索引再取数：

# 先拆分原数据集的行索引
train_idx, test_idx = train_test_split(X.index, test_size=0.20, random_state=99)
# 用索引取对应行数据，自动保留原有列名和索引
X_train, X_test = X.loc[train_idx], X.loc[test_idx]
y_train, y_test = y.loc[train_idx], y.loc[test_idx]

如果不需要和原数据集对齐，可直接重置为从0开始的连续索引：

X_train = X_train.reset_index(drop=True)
X_test = X_test.reset_index(drop=True)
y_train = pd.Series(y_train).reset_index(drop=True)
y_test = pd.Series(y_test).reset_index(drop=True)

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