筛选pandas DataFrame按日期分组后的前N条与后N条行数据

实现思路 #

按date字段分组后，对每个分组的val列做自定义方向排序，取排序后分组的前N行和后N行，合并去重后即为结果，可灵活调整N值适配不同场景。

通用代码实现 #

首先是示例数据构造：

import pandas as pd
df = pd.DataFrame({'date':['11-10','11-10','11-10','12-10','12-10','12-10'],
                   'id': [1, 1, 2, 1, 1, 2],
                   'val':[20, 30, 40, 50, 25, 35] })

方法1：逻辑直观版（适合中小数据量）

通过groupby+apply实现，易修改调整逻辑：

def filter_group_top_bottom(df, group_col:str, sort_col:str, n:int, ascending:bool=True):
    def single_group_process(g):
        sorted_g = g.sort_values(sort_col, ascending=ascending)
        # 拼接前N、后N后去重，避免排序后中间数据重叠重复保留
        return pd.concat([sorted_g.head(n), sorted_g.tail(n)]).drop_duplicates()
    return df.groupby(group_col, group_keys=False).apply(single_group_process).reset_index(drop=True)

方法2：高性能版（适合百万级以上大数据量）

用rank方法避免逐组迭代，性能提升明显：

def filter_group_top_bottom_fast(df, group_col:str, sort_col:str, n:int, ascending:bool=True):
    df['group_rank'] = df.groupby(group_col)[sort_col].rank(ascending=ascending, method='first')
    group_size = df.groupby(group_col)[sort_col].transform('count')
    # 同时满足前N或后N条件
    filter_df = df[(df['group_rank'] <= n) | (df['group_rank'] >= group_size - n + 1)]
    return filter_df.drop('group_rank', axis=1).reset_index(drop=True)

示例测试（匹配你给出的N=2预期结果） #

你给出的示例需要删除每组val最小的1条，保留最大的2条，直接调用方法即可：

N = 2
# ascending=False 表示按val从大到小排序，取前2条刚好符合你的预期
result = filter_group_top_bottom(df, group_col='date', sort_col='val', n=N, ascending=False)

输出结果完全符合预期：

如果需要同时保留每组最小的N条和最大的N条，只需将ascending参数改为True即可。

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