有条件地将数据拆分为训练和测试（Pandas）

Question

我有一个使用 Python 执行预测任务的代码。 任务是预测一家公司从 2015 年到 2019 年不同年份的销售额。

我想将数据拆分为训练集和测试集。

但问题是，我想用2015年到2018年的数据训练模型，在2019年的数据上测试模型。

我怎样才能使用 train_test_split、ShuffleSplit、有条件地拆分数据？

X_train = df.iloc[train_index]
X_test = df.iloc[test_index]
y_train = X_train.Sales
y_test = X_test.Sales

Answer 1

由于您在一开始就遇到了条件，因此您将失去使用机器学习预处理中使用的改组方法的好处。 因此，我建议不要在这种情况下执行训练测试拆分（我假设结果有偏差）。 不过，如果您需要这样做，请尝试：

train = your_data[your_data['year_column'] < 2019]
test = your_data[your_data['year_column'] == 2019]

X_train = train.loc[:, train.columns != 'column_of_interest']
y_train = train['column_of_interest']
X_test = test.loc[:, test.columns != 'column_of_interest']
y_train = test['column_of_interest']