如何用Pandas的groupby+transform为实验组计算RMSE并添加列？

需求：为DataFrame中的不同实验组计算均方根误差（RMSE），并将对应组的RMSE值新增为一列，每组内所有行的RMSE值保持一致。

示例输入DataFrame #

import pandas as pd
from scipy import metrics

df = pd.DataFrame({
    'experiment': ['A', 'A', 'B', 'B'],
    'reaction': [1, 2, 1, 2],
    'result': [1.5, 1.9, 2.1, 2.4],
    'target': [2.0, 1.6, 2.0, 1.6]
})

期望输出 #

experiment  reaction  result  target   rmse
0          A         1     1.5     2.0  0.41
1          A         2     1.9     1.6  0.41
2          B         1     2.1     2.0  0.57
3          B         2     2.4     1.6  0.57

方案一：groupby + transform #

利用transform方法对每个分组执行计算，并自动返回与原数据长度匹配的结果，直接新增列即可：

# 定义组内RMSE计算函数
def calculate_rmse(group):
    return metrics.mean_squared_error(group['result'], group['target'], squared=False)

# 新增rmse列并保留两位小数
df['rmse'] = df.groupby('experiment').transform(calculate_rmse).round(2)

transform会将每个组的RMSE值重复对应组的行数，直接填充到新列，无需额外合并操作。

方案二：groupby + apply + merge #

先通过apply计算各组RMSE，得到单独的结果表，再与原DataFrame合并：

# 分组计算RMSE，转换为带列名的DataFrame
rmse_df = df.groupby('experiment').apply(
    lambda x: metrics.mean_squared_error(x['result'], x['target'], squared=False)
).reset_index(name='rmse')

# 合并到原DataFrame并保留两位小数
df = df.merge(rmse_df.round(2), on='experiment')

该方法逻辑清晰，适合需要单独保存各组RMSE统计结果的场景。

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