如何从Pandas的groupby对象中提取每组的指定行？

问题场景 #

现有按ID分组的Pandas DataFrame分组对象，数据示例如下：

sl.no   ID      hour
50834   1       6
50833   1       6
50832   1       6
19188   2       9
19187   2       9
19186   2       9
13380   3       294
13379   3       294
13378   3       294

可通过以下代码创建该分组对象：

import pandas as pd

data = {
    'sl.no': [50834, 50833, 50832, 19188, 19187, 19186, 13380, 13379, 13378],
    'ID': [1, 1, 1, 2, 2, 2, 3, 3, 3],
    'hour': [6, 6, 6, 9, 9, 9, 294, 294, 294]
}
df = pd.DataFrame(data)

dfnew = df.sort_values(['sl.no'], ascending=False).groupby('ID').head(3)

需求为提取每组的第n行（例如提取第2行），期望输出如下：

sl.no   ID      hour
50833   1       6
19187   2       9
13379   3       294

此前尝试用iloc(2)无法实现，通过head(2)与head(1)求差的方式实现，现寻求更简便直接的方法。

解决方案 #

方法1：使用groupby.nth()（推荐） #

Pandas的groupby.nth()方法可以直接提取每组指定位置的行，注意索引是从0开始计数的。比如要提取每组的第2行（对应索引1），代码如下：

# 提取每组第2行（0-based索引为1）
result = dfnew.groupby('ID').nth(1)
print(result)

运行后输出与期望一致：

sl.no  hour
ID
1      50833     6
2      19187     9
3      13379   294

nth()还支持更多场景：

• 提取多个位置：nth([0,2])可提取每组第1和第3行
• 提取倒数第n行：nth(-1)提取每组最后一行

方法2：使用apply() + iloc #

也可以通过apply()结合iloc实现，代码如下：

result = dfnew.groupby('ID').apply(lambda x: x.iloc[1]).reset_index(drop=True)
print(result)

输出结果：

sl.no  ID  hour
0  50833   1     6
1  19187   2     9
2  13379   3   294

不过这种方法性能略低于nth()，数据量较大时推荐用第一种方法。

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