如何按组内指定SampleSize随机抽取样本并添加标记列？

你的问题出在自定义函数返回的是选中ID的子集，无法直接和原DataFrame对应生成标记列。可以通过调整函数逻辑，直接在分组内完成标记，具体方案如下：

方法一：修改自定义函数配合groupby.apply #

import pandas as pd
import numpy as np

# 原数据集
df = pd.DataFrame({
    'ID': [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
    'Group': [1, 1, 1, 1, 1, 1, 2, 2, 2, 2],
    'SampleSize': [4, 4, 4, 4, 4, 4, 1, 1, 1, 1]
})

def mark_selected_samples(group):
    sample_size = group['SampleSize'].iloc[0]
    # 初始化Sample列为0
    group['Sample'] = 0
    # 从分组索引中随机选指定数量的行，标记为1
    selected_indices = np.random.choice(group.index, size=sample_size, replace=False)
    group.loc[selected_indices, 'Sample'] = 1
    return group

# 按Group分组处理，拼接回完整结果
result_df = df.groupby('Group', group_keys=False).apply(mark_selected_samples)
print(result_df)

关键说明 #

• 函数直接在传入的分组DataFrame上新增标记列，保留所有行数据
• 用分组的索引而非ID做随机选择，避免ID重复时的错误
• group_keys=False避免生成额外的分组索引层级

方法二：更简洁的transform写法（推荐） #

如果追求代码精简，可借助groupby.transform直接映射标记结果：

def get_sample_mask(x):
    sample_size = df.loc[x.index, 'SampleSize'].iloc[0]
    selected_ids = np.random.choice(x, size=sample_size, replace=False)
    return x.isin(selected_ids).astype(int)

df['Sample'] = df.groupby('Group')['ID'].transform(get_sample_mask)

关键说明 #

• transform会自动将分组内的处理结果映射回原DataFrame的对应行
• astype(int)将布尔值（是否被选中）转为1/0的整数标记

原方法失效原因 #

你之前的函数仅返回包含选中ID的小DataFrame，groupby.apply后会丢失未被选中的行，无法和原数据对应生成完整的标记列。调整后让函数返回完整分组数据并新增标记列，即可解决问题。

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