sklearn拆分数据集后导出训练/测试集CSV文件的实现方法

pd.DataFrame(X_train, Y_train)写法存在参数传递错误：pandas中DataFrame构造函数的第二个位置参数是行索引参数index，你传入的Y_train会被当做行索引使用，不会和X_train做列拼接，索引对齐失败就会出现所有特征列全为NaN的情况。

使用pd.concat()做列方向的横向拼接，指定axis=1即可把特征和对应标签按行对齐合并，之后直接调用to_csv()导出即可，完整代码如下：

# 保留你原有的数据集拆分逻辑
# X_train, X_test, Y_train, Y_test = train_test_split(X,Y,test_size = .3, random_state = 50)

# 按列拼接训练集、测试集的特征与标签
training_set = pd.concat([X_train, Y_train], axis=1)
test_set = pd.concat([X_test, Y_test], axis=1)

# 导出为CSV文件，index=False 用于取消写入pandas自动生成的行索引
training_set.to_csv("training_set.csv", index=False)
test_set.to_csv("test_set.csv", index=False)

如果拆分前你没有给标签列设置列名（比如Y是一维Series/数组没有对应列名），导出前可以手动重命名最后一列（即标签列），以iris数据集的标签列variety为例：

training_set = training_set.rename(columns={training_set.columns[-1]: "variety"})
test_set = test_set.rename(columns={test_set.columns[-1]: "variety"})

不要使用pd.DataFrame()直接传两个数据集做拼接，这个构造方法不支持直接传入两个同维度数据集做横向合并。

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