如何在Pandas中按分组标签动态提取每组前N行？

需求：现有一个DataFrame，需按label列分组，根据每个分组标签对应的函数返回的N值，提取每组的前N行。固定N值时可直接用df.groupby('label').head(n)，但N随标签变化时该方法不适用，且因对性能敏感，需避免遍历df['label'].unique()拼接数据的方案。

示例数据 #

label   values
0  apple       7
1  apple       5
2  apple       4
3    car       9
4    car       6
5    dog       5
6    dog       3
7    dog       2
8    dog       1

示例代码 #

import pandas as pd

df = pd.DataFrame({
    'label': ['apple', 'apple', 'apple', 'car', 'car', 'dog', 'dog', 'dog', 'dog'],
    'values': [7, 5, 4, 9, 6, 5, 3, 2, 1]
})

def depth(label):
    if label == 'apple': return 1
    elif label == 'car': return 2
    elif label == 'dog': return 3

高效解决方案 #

利用Pandas的矢量化操作和分组计数功能，避免循环遍历，实现高效筛选：

方法1：直接用apply映射N值 #

# 为每行添加对应标签的N值
df['required_n'] = df['label'].apply(depth)
# 按label分组后，计算组内行索引（从0开始），筛选索引小于required_n的行
result = df[df.groupby('label').cumcount() < df['required_n']].drop('required_n', axis=1)
print(result)

方法2：用字典映射替代apply（性能更优） #

将函数转换为字典，用map实现更快的矢量化映射：

depth_map = {
    'apple': 1,
    'car': 2,
    'dog': 3
}

df['required_n'] = df['label'].map(depth_map)
result = df[df.groupby('label').cumcount() < df['required_n']].drop('required_n', axis=1)

期望输出 #

label  values
0  apple       7
3    car       9
4    car       6
5    dog       5
6    dog       3
7    dog       2

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