基于切片更新pandas数据帧

Question

我有一个数据帧，希望使用sklearn.model_selection.train_test_split函数将其拆分为"训练"one_answers"测试"数据集。此函数返回原始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了。但我很好奇我怎么能以不同的方式做到这一点。我想，我可以直接编辑df，而不是编辑train_df和test_df并再次添加它们，但由于我不太熟悉.loc和.iloc的工作原理，我很难理解这将如何工作。

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

df["Dataset"] = "train" WHERE index in train_df.index.values
df["Dataset"] = "test" WHERE index in test_df.index.values

Answer 1

如果您不想在model_selection.train_test_split()调用中copy您的DataFrame，您可以使用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来获得适应性更强的解决方案：

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