如何在Pandas中按因子Y的每个层级随机抽取最多n行数据？

可以通过groupby结合sample方法轻松实现这个需求，核心逻辑是对每个因子层级，抽取该层级样本数与n的较小值：

步骤示例 #

1. 构造示例数据 #

先创建一个符合需求的DataFrame（实际使用时替换成你的数据即可）：

import pandas as pd
import numpy as np

# 设置随机种子确保结果可复现
np.random.seed(42)
# 生成样本ID和因子Y的层级数据
sample_ids = [f"sample_{i}" for i in range(100)]
Y = np.random.choice(['X', 'Z', 'Y'], size=100, p=[0.03, 0.3, 0.67])
# 构建以sample_id为行索引的DataFrame
df = pd.DataFrame({'Y': Y}, index=sample_ids)

2. 执行分层随机抽样 #

n = 10  # 每组最多抽取的行数
# 按Y分组后，每个组抽取min(组内样本数, n)行
sampled_df = df.groupby('Y').apply(
    lambda group: group.sample(n=min(len(group), n), random_state=42)
).reset_index(drop=True)

3. 验证结果 #

可以查看每个层级最终抽取的样本数量：

print(sampled_df['Y'].value_counts())

代码说明 #

• groupby('Y')：将DataFrame按因子Y的不同层级分组
• group.sample(n=min(len(group), n))：对每个分组，判断组内样本数：如果小于等于n则全量保留，否则随机抽取n行
• random_state：可选参数，设置后可以固定随机抽样的结果，方便调试和复现
• reset_index(drop=True)：移除分组后产生的多级索引，生成整洁的结果DataFrame

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