按分组获取每组指定末尾N条记录（DataFrame解决方案）

输入数据 #

import pandas as pd

df = pd.DataFrame({
    'id': ["a","a","a","a","b","b","b"],
    'tail_num': [2,2,2,2,1,1,1],
    'value': [1,2,3,4,5,6,7]
})

需求说明 #

按id字段分组，每个分组需保留末尾N条记录，其中N为该分组tail_num列的统一值（如id='a'的分组取末尾2条，id='b'的分组取末尾1条）。

期望输出 #

expected_df = pd.DataFrame({
    'id': ["a","a","b"],
    'tail_num': [2,2,1],
    'value': [3,4,7]
})

解决方案 #

方法1：groupby.apply 直观实现

适合小数据集，逻辑清晰易懂：

# 按id分组，每个分组取末尾N条（N为分组内tail_num的第一个值）
result = df.groupby('id').apply(
    lambda x: x.tail(x['tail_num'].iloc[0])
).reset_index(drop=True)

方法2：向量化操作 高效实现

适合大数据集，避免逐组循环的性能损耗：

# 生成每个分组内的倒序位置标记（最后一条为0，倒数第二条为1，以此类推）
df['reverse_rank'] = df.groupby('id').cumcount(ascending=False)
# 筛选位置标记小于tail_num的记录，即保留末尾N条
result = df[df['reverse_rank'] < df['tail_num']].drop('reverse_rank', axis=1).reset_index(drop=True)

结果验证 #

执行以下代码可确认结果与期望完全一致：

print(result.equals(expected_df))  # 输出: True

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