我想在一个Pandas数据框上进行测试、训练和验证，但我不想生成新的数据框。相反，我想添加一个名为’Split’的新列，其中Split = ['train','valid','test']。我希望'train'、'valid'和'test'分别随机分布在64%、16%和20%的行中。

我知道scikit-learn的train_test_split，但我再次强调，我不想要新的框架。所以我可以尝试：

from sklearn.model_selection import train_test_splittrain, test = train_test_split(df, test_size=0.2)

但我只想要一个名为’Split’的列，其值为train、valid和test作为标签。这是为了机器学习目的，所以我希望确保分割是完全随机的。

有谁知道这可能实现吗？

回答：

这里有一种方法，使用建议的numpy.random.choice：

import pandas as pdimport numpy as np# 设置一个小示例data = np.ones(shape=(100, 3))df = pd.DataFrame(data, columns=['x1', 'x2', 'y'])df['split'] = pd.NA# 分割split = ['train', 'valid', 'test']df['split'] = df['split'].apply(lambda x: np.random.choice(split, p=[0.64, 0.16, 0.20]))# 验证df['split'].value_counts()

在一次运行中，结果如下

train    64valid    19test     17Name: split, dtype: int64