Question

我有一个数据框，希望使用sklearn.model_selection.train_test_split函数将其分为“训练”和“测试”数据集。此函数返回原始DataFrame的两个切片。但是，我需要将其放在单个DataFrame中，并带有一个标识条目类型的列条目。我可以编写一个代替此功能的函数，但是使用sklearn函数既方便又可靠。

我当前的方法如下：

import pandas as pd
import numpy as np
from sklearn import model_selection

dates = pd.date_range('20130101',periods=10)
df = pd.DataFrame(np.random.randn(10,4),index=dates,columns=list('ABCD')).reset_index()

split = [0.8, 0.2]
split_seed = 123

train_df, test_df = model_selection.train_test_split(df, train_size = split[0], test_size = split[1], random_state=split_seed)

train_df["Dataset"] = "train"
test_df["Dataset"] = "test"

final_df = train_df.append(test_df)

这很好用，但是由于我正在修改复制的切片而不是原始的df对象，因此会导致警告：

A value is trying to be set on a copy of a slice from a DataFrame.
Try using .loc[row_indexer,col_indexer] = value instead

这并不重要，因为此后不再使用原始DataFrame。但是我很好奇如何才能做到这一点。我假设不用编辑train_df和test_df并再次附加它们，我可以直接编辑df，但是由于我对.loc和{ {1}}工作正常，我正在努力查看如何工作。

说明我正在寻找的伪代码如下：

.iloc

Answer 1

如果您不想在copy通话中DataFrame model_selection.train_test_split()，则可以使用loc：

df.loc[train_df.index, 'Dataset'] = 'train'
df.loc[test_df.index, 'Dataset'] = 'test'

Answer 2

一种方法是使用np.where添加一个以布尔条件为条件的序列：

df['Dataset'] = np.where(df.index.isin(train_df.index.values), 'train', 'test')

当然，这假定train_df中没有包含的索引必须中存在。

或使用test_df获得更具适应性的解决方案：

np.select

根据切片更新熊猫数据框

2 个答案: