如何对DataFrame按diet分组过滤，每组最多保留5条记录？

我有一个包含diet因子列的DataFrame，该列的取值为"herbivores"、"carnivores"和"omnivores"，对应分组的记录数分别为3条、6条和8条。需要对DataFrame做如下过滤处理：

• 记录数少于5的分组保持原样
• 记录数超过5的分组随机筛选至5条记录

原DataFrame示例 #

diet       factor2
1 herbivore     a
2 herbivore     a
3 herbivore     a
4 carnivore     a
5 carnivore     a
6 carnivore     a
7 carnivore     a
8 carnivore     a
9 carnivore     a
10 omnivore     a
11 omnivore     a
12 omnivore     a
13 omnivore     a
14 omnivore     a
15 omnivore     a
16 omnivore     a
17 omnivore     a

期望过滤后的DataFrame #

diet    factor2
1 herbivore     a
2 herbivore     a
3 herbivore     a
4 carnivore     a
5 carnivore     a
6 carnivore     a
7 carnivore     a
8 carnivore     a
9 omnivore     a
10 omnivore     a
11 omnivore     a
12 omnivore     a
13 omnivore     a

解决方案 #

用pandas的groupby结合apply方法即可实现需求，核心逻辑是对每个分组判断记录数，按需保留或采样：

import pandas as pd

# 构造示例DataFrame
data = {
    'diet': ['herbivore']*3 + ['carnivore']*6 + ['omnivore']*8,
    'factor2': ['a']*17
}
df = pd.DataFrame(data)

# 分组过滤处理
filtered_df = df.groupby('diet', group_keys=False).apply(
    lambda x: x if len(x) <= 5 else x.sample(n=5, random_state=42)
).reset_index(drop=True)

print(filtered_df)

代码说明 #

• groupby('diet')按diet列对数据分组
• apply内的匿名函数对每个分组做判断：长度≤5则直接保留全部数据，超过5则用sample(n=5)随机选取5条
• random_state=42用于固定采样结果，保证多次运行输出一致，不需要固定结果可删除该参数
• reset_index(drop=True)重置索引，避免分组后索引混乱

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