按列a分组后，提取列b最大值对应的列c值并新增列赋值

需求说明 #

在按列a分组的DataFrame中，根据列b（日期类型）的最大值所在行，提取对应列c的值，并将该值赋值到新增列d中。当前数据集包含超过16000行数据。

示例数据 #

import pandas as pd
import numpy as np

df = pd.DataFrame({'a': ['a','a','a','b','b','b','c','c','c'],
                   'b': ['2008-11-01', '2022-07-01', '2017-02-01', '2017-02-01', '2018-02-01', '2008-11-01', '2014-11-01', '2008-11-01', '2022-07-01'],
                   'c': [np.NaN,6,8,2,1,np.NaN,6,np.NaN,7]})
df['b'] = pd.to_datetime(df['b'])

期望结果 #

a          b    c    d
0   a 2008-11-01  NaN  8.0
1   a 2022-07-01  6.0  8.0
2   a 2017-02-01  8.0  8.0
3   b 2017-02-01  2.0  1.0
4   b 2018-02-01  1.0  1.0
5   b 2008-11-01  NaN  1.0
6   c 2014-11-01  6.0  7.0
7   c 2008-11-01  NaN  7.0
8   c 2022-07-01  7.0  7.0

尝试过的方法及报错 #

• 执行df['d'] = df.groupby('a').b.astype(np.int64).max()[df.c].reset_index()时，提示分组后的SeriesGroupBy无法转换类型
• 提前将列b转为int类型后，报错提示列c的值不在索引中
• 尝试分组后用where函数时，因未指定替代值报错

可行解决方案 #

方法一：使用groupby.transform结合idxmax #

该方法直接在原DataFrame上计算，无需额外合并，效率较高，适合大数据集：

# 对每个分组，找到b列最大值对应的行索引，提取对应c值后广播到整个分组
df['d'] = df.groupby('a')['c'].transform(lambda x: x.loc[x.index[df.loc[x.index, 'b'].idxmax()]])

方法二：先分组获取最大值对应数据，再合并到原DataFrame #

逻辑直观，便于理解分组操作：

# 1. 按a分组，筛选每组b列最大的行，提取a和c列并重命名c为d
max_b_rows = df.groupby('a').apply(lambda x: x[x['b'] == x['b'].max()][['a', 'c']]).rename(columns={'c': 'd'}).reset_index(drop=True)
# 2. 将结果合并到原DataFrame
df = df.merge(max_b_rows, on='a', how='left')

方法三：使用groupby.idxmax定位行后映射 #

通过字典映射实现，执行效率优秀：

# 获取每个分组中b列最大值对应的行索引
max_indices = df.groupby('a')['b'].idxmax()
# 创建a到对应c值的映射字典
a_to_d = df.loc[max_indices, ['a', 'c']].set_index('a')['c'].to_dict()
# 映射到原DataFrame生成d列
df['d'] = df['a'].map(a_to_d)

以上三种方法均可实现需求，其中方法一和方法三更适合处理16000+行的数据集。

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