Pandas多分组提取首值异常：groupby未识别index分组如何修正？

需求说明 #

需要在Pandas DataFrame中创建desired_output列，取值为每个ID下每个index分组对应的lumpsum字段的第一行值。

示例数据与初始化代码 #

import pandas as pd

data = [
    [1, 12334, 1, 12334],
    [1, 12334, 1, 12334],
    [1, 12334, 1, 12334],
    [1, 12334, 1, 12334],
    [1, 34567, 1, 12334],
    [1, 34567, 1, 12334],
    [2, 45788, 1, 45788],
    [2, 45788, 2, 45788],
    [2, 23467, 2, 45788],
    [2, 5678, 3, 5678],
    [2, 4567, 3, 5678],
    [3, 56832, 1, 56832],
    [3, 43456, 1, 56832],
    [3, 2378, 2, 2378],
    [4, 6754, 1, 6754],
    [4, 3456, 2, 3456]
]

columns = ['ID', 'lumpsum', 'index', 'desired_output']
df = pd.DataFrame(data, columns=columns)
print(df)

问题分析 #

你尝试的代码：

df['test']=df.groupby('ID', 'index')['lumpsum'].transform('first')

执行后仅按ID分组提取lumpsum的第一行值，完全忽略index分组的原因是：Pandas的groupby方法接收多列分组时，需要将列名放在一个列表中传递。你单独传入两个参数时，第二个参数'index'会被识别为axis参数（默认值为0），因此实际只按ID列进行了分组。

修正方案 #

将分组列名放入列表中，修改后的代码如下：

df['test'] = df.groupby(['ID', 'index'])['lumpsum'].transform('first')

这样就能实现按ID和index的组合分组，提取每个分组内lumpsum字段的第一行值，并将结果映射回原DataFrame的对应行。验证后，test列的取值会和示例中的desired_output列完全一致。

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