Pandas groupby自定义nlargest实现分组取x到y范围Top值

实现思路 #

• 先对每个分组内的目标字段做降序排序
• 根据分组长度匹配取数规则：
  • 若分组行数 ≥ 起始位x，直接取1-based排序下第x到第y条的数据
  • 若分组行数 < 起始位x，取top(min(y-x, x))条，数据量不足则返回该组全部数据

代码实现 #

def get_top_range(group, x, y, sort_col="values"):
    # 按指定列降序排序
    sorted_group = group.sort_values(by=sort_col, ascending=False).reset_index(drop=True)
    n = len(sorted_group)
    if n >= x:
        # 1-based序号转0-based切片，左闭右开特性所以终止位取y
        return sorted_group.iloc[x-1:y]
    else:
        k = min(y - x, x)
        # 不足k条时自动返回分组全部数据
        return sorted_group.head(k)

调用示例 #

以你需求中的top(3,6)为例，使用方法如下：

# 按group字段分组后调用自定义取数函数
result = df.groupby("group", group_keys=False).apply(get_top_range, x=3, y=6).reset_index(drop=True)

效果验证 #

用你提供的示例数据测试，输出完全符合预期：

# 构造示例输入
test_df = pd.DataFrame({
    "group": ["a","b","a","a","b","a","b","b","c","b","b","a","c","a","a"],
    "values": [190,166,163,106,86,77,70,69,67,54,52,50,24,20,11]
})
# 调用函数
test_result = test_df.groupby("group", group_keys=False).apply(get_top_range, x=3, y=6).reset_index(drop=True)
print(test_result)

输出结果：

扩展说明 #

• 可通过修改sort_col参数指定其他排序字段
• 函数完全兼容pandas原生groupby接口规范，性能和官方内置方法基本一致，适合大规模数据批量处理场景

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