如何对pandas DataFrame分组后按自定义字典指定数量调用head/tail取行

方法1：groupby + apply（易读易维护） #

该写法逻辑直观，适合中小数据集：

import pandas as pd

# 测试输入数据
df = pd.DataFrame({
    "group": ["A","A","A","B","B","B","B"],
    "value": [0,1,2,3,4,5,6]
})
slice_per_group = {"A": 1, "B": 3}

# 核心实现逻辑
result = df.groupby('group', group_keys=False).apply(
    lambda g: g.head(slice_per_group.get(g.name, 0))
)
print(result)

参数说明：

• group_keys=False 避免生成额外的分组层级索引
• slice_per_group.get(g.name, 0) 从字典中读取当前分组需要取的行数，第二个参数为默认值，未在字典中配置的分组默认返回0行，若需要返回未配置分组的全部数据可将默认值改成len(g)

方法2：cumcount + map（高性能版） #

大数据量下性能更优，避免apply逐组遍历的开销：

# 生成分组内从0开始的行序号
df['group_inner_idx'] = df.groupby('group').cumcount()
# 映射分组对应的行数，过滤符合要求的行
result = df[df['group_inner_idx'] < df['group'].map(slice_per_group)].drop('group_inner_idx', axis=1)
print(result)

扩展说明 #

如果需要实现按分组取末尾N行（对应tail功能），只需做小幅调整：

• 方法1：将head替换为tail
• 方法2：将cumcount()替换为cumcount(ascending=False)

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