sklearn train_test_split指定最后1行为测试集其余为训练集

你这个固定取最后1行作为测试集的场景，不需要使用train_test_split——这个函数的作用是随机拆分数据集，没法精准保证最后一行被分到测试集，直接用pandas的位置切片就能实现需求，代码更简洁逻辑也更准确：

import pandas as pd

# 读取csv文件
df = pd.read_csv('Book5.csv', names=['Amiability', 'Email'])

# 拆分数据集：除最后一行外全部作为训练集，最后一行作为测试集
train_set = df.iloc[:-1, :]
test_set = df.iloc[-1:, :]

# 拆分特征和标签
x_train = train_set['Amiability']
y_train = train_set['Email']
x_test = test_set['Amiability']
y_test = test_set['Email']

代码说明 #

• iloc是pandas按位置索引取值的方法，第一个参数控制行选择，第二个参数控制列选择
• [:-1]表示选取从第1行到倒数第二行的所有内容，刚好对应你需要的0~330行（共330条）训练数据
• [-1:]表示只选取最后1行作为测试集，保留:是为了让返回结果保持DataFrame/Series的结构，避免后续传入sklearn模型预测时出现维度不匹配的报错
• 如果你使用的模型要求输入特征为二维数组，可以把特征选取的代码改成x_train = train_set[['Amiability']]、x_test = test_set[['Amiability']]即可，双中括号取列会返回二维结构的DataFrame。

额外提示：你原本写的train_test_split逻辑是按20%比例随机拆分，在你当前331行数据的情况下测试集约为66条数据，和你要1条测试集的需求完全不符，不需要保留这段代码。

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