如何为Pandas DataFrame新增列，按比例随机标注train/valid/test拆分标签

下面分两种常用实现路径，都可以直接在原DataFrame上新增Split列，不会生成额外的子DataFrame：

方法1：用Pandas+NumPy直接划分（最简洁） #

仅依赖你已经在用的pandas和numpy库，一次随机抽样完成划分，占比精准：

import pandas as pd
import numpy as np

# 设置随机种子保证实验可复现，不需要可删除该行
np.random.seed(42)

# 生成0-1区间的均匀随机数
random_arr = np.random.rand(len(df))
# 按阈值划分三类标签
df['Split'] = pd.cut(
    random_arr,
    bins=[0, 0.64, 0.8, 1],
    labels=['train', 'valid', 'test'],
    include_lowest=True
)

方法2：基于scikit-learn的train_test_split实现（适配分层抽样需求） #

如果你已经熟悉sklearn的拆分逻辑，或者需要做分层抽样保证标签分布一致，可以直接操作索引打标，不需要拆分新DataFrame：

from sklearn.model_selection import train_test_split

# 设置随机种子保证可复现
RANDOM_STATE = 42

# 第一步：拆分出占比20%的测试集索引
train_valid_idx, test_idx = train_test_split(
    df.index,
    test_size=0.2,
    random_state=RANDOM_STATE,
    # 如需分层抽样取消注释下行，替换为你的标签列名
    # stratify=df['label']
)

# 第二步：在剩余80%的样本中，拆分出占总样本16%的验证集（即剩余样本的20%）
train_idx, valid_idx = train_test_split(
    train_valid_idx,
    test_size=0.2,
    random_state=RANDOM_STATE,
    # 如需分层抽样同步取消注释下行
    # stratify=df.loc[train_valid_idx, 'label']
)

# 给对应索引行打标签
df['Split'] = 'train'
df.loc[valid_idx, 'Split'] = 'valid'
df.loc[test_idx, 'Split'] = 'test'

可以执行以下代码验证拆分占比是否符合预期：

print(df['Split'].value_counts(normalize=True).round(4)*100)

正常输出会接近如下结果：

train    64.0
test     20.0
valid    16.0
Name: Split, dtype: float64

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