如何实现Pandas DataFrame分组的相等概率抽样（非固定行数）

假设你的DataFrame结构如下：

import pandas as pd

df = pd.DataFrame({
    'ID': ['a', 'a', 'b', 'c', 'c', 'c'],
    'Value': [2, 4, 6, 8, 10, 12]
})

方法1：先等概率抽取分组，再处理组内行 #

这种方式先从所有ID分组中以相等概率选中目标分组，再提取这些分组的全部或部分行，确保每个分组被选中的概率一致。

示例：随机抽取指定数量的分组 #

# 获取所有唯一ID分组
unique_ids = df['ID'].unique()
# 等概率抽样1个分组（replace=True允许重复抽选，False则不允许）
sampled_ids = pd.Series(unique_ids).sample(n=1, replace=False).tolist()
# 筛选出对应分组的所有行
result = df[df['ID'].isin(sampled_ids)]

示例：给每个分组设置固定被选中概率，选中后组内抽样 #

def sample_group(group):
    # 每个分组有50%的概率被选中，可自行调整概率值
    if pd.np.random.rand() < 0.5:
        return group.sample(n=1)  # 组内随机抽取1行
    return pd.DataFrame()  # 未选中则返回空表

# 按ID分组应用抽样逻辑
result = df.groupby('ID').apply(sample_group).reset_index(drop=True)

方法2：按分组大小设置抽样权重，保证组级概率均等 #

如果需要从全量行中抽样，但让每个分组的整体被抽样概率相同（避免大组占比过高），可以给组内每行设置与组大小成反比的权重，确保每个组的总权重一致：

# 计算每个分组的行数
group_sizes = df.groupby('ID').size()
# 给每行分配权重：1/对应分组的行数，使每个组总权重为1
df['weight'] = df['ID'].map(lambda x: 1 / group_sizes[x])
# 按权重抽样，这里抽3行，每个分组的被抽样概率均等
result = df.sample(n=3, weights='weight', replace=False)

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