如何从DataFrame各类别中抽取最多1000条样本？

需求说明 #

从DataFrame的每个类别（以class列的取值为例）中抽取样本，规则为：每个类别最多保留1000行数据；若该类别总行数不足1000，则直接保留全部行。

示例数据 #

import pandas as pd

data = {
    'Date': ['2023-10-20', '2023-10-21', '2023-10-22', '2023-10-23', '2023-10-24'],
    'class': [4, 1, 1, 2, 3],
    'other_col1': [5, 6, 3, 1, 4],
    'other_col2': [15, 10, 72, 6, 8]
}

df = pd.DataFrame(data)

基础循环实现（已尝试） #

sample_size = 1000
sampled_data = []

for class_label, group in df.groupby('class'):
    if len(group) >= sample_size:
        sampled_data.append(group.sample(sample_size))
    else:
        sampled_data.append(group)

sampled_df = pd.concat(sampled_data)

更简洁的实现方式 #

利用groupby().apply()结合sample()方法，一行代码完成需求，无需手动拼接结果：

sample_size = 1000
sampled_df = df.groupby('class').apply(
    lambda group: group.sample(n=min(len(group), sample_size), random_state=42)
).reset_index(drop=True)

说明 #

• lambda group: group.sample(n=min(len(group), sample_size))：对每个分组自动判断抽样数量，取sample_size和分组行数的较小值，省去手动分支判断
• random_state=42：可选参数，用于固定抽样结果的随机性，保证可复现；不需要的话可直接删除
• reset_index(drop=True)：重置结果索引，避免分组后遗留的多级索引干扰后续操作

内容的提问来源于stack exchange，提问作者Anna Andreeva Rogotulka