Pandas中groupby.agg结合nth()提取分组列指定值失败问题排查

你的代码主要存在以下几个问题：

1. lambda函数列名混乱：用字典格式agg时，lambda函数会被自动命名为<lambda>，无法生成second/third这类清晰的自定义列名。
2. 未处理分组长度不足的场景：当分组元素数量少于指定的nth位置时，会返回NaN，原代码未做针对性处理。
3. 多层列结构不符合预期：聚合结果会生成双层索引列（外层为A，内层为聚合方法名），和你想要的单层列结构不符。

正确实现方式一：使用命名聚合（推荐） #

Pandas 0.25+支持命名聚合，可以直接为每个聚合规则指定列名，同时通过判断分组长度避免无效的NaN：

import pandas as pd

df = pd.DataFrame({'index1':[1,2,2,2,3,3], 'A':[1,1,2,3,1,4]})

# 命名聚合，自定义列名并处理分组长度不足的情况
result = df.groupby('index1').agg(
    first=('A', 'first'),
    second=('A', lambda x: x.nth(1) if len(x) >= 2 else pd.NA),
    third=('A', lambda x: x.nth(2) if len(x) >= 3 else pd.NA),
    fourth=('A', lambda x: x.nth(3) if len(x) >= 4 else pd.NA),
    last=('A', 'last')
).reset_index()

# 将NaN替换为空字符串，匹配预期输出格式
result = result.fillna('')
print(result)

输出结果：

index1 first second third fourth last
0       1     1                      1
1       2     1      2     3          3
2       3     1      4                4

正确实现方式二：先编号再转宽表 #

通过给每个分组内的元素按顺序编号，再用pivot转成宽表，最后补充last列：

import pandas as pd

df = pd.DataFrame({'index1':[1,2,2,2,3,3], 'A':[1,1,2,3,1,4]})

# 为每个分组内的元素分配顺序编号
df['seq'] = df.groupby('index1').cumcount()

# 转成宽表并重命名列
pivot_df = df.pivot(index='index1', columns='seq', values='A').rename(
    columns={0: 'first', 1: 'second', 2: 'third', 3: 'fourth'}
)

# 添加last列
pivot_df['last'] = df.groupby('index1')['A'].last()

# 重置索引并替换空值
result = pivot_df.reset_index().fillna('')
print(result)

这种方法逻辑更直观，适合需要扩展更多位置提取的场景。

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