Pandas按组计算相关性并新增列：为何均值方法失效？

现有如下数据框：

import pandas as pd
sample_df = pd.DataFrame({'id': ['A', 'A', 'A', 'A', 'B', 'B', 'B', 'B'],
                          'values_1':[2,4,6,8,12,13,13,17],
                          'values_2':[3,6,7,9,3,2,2,5]})

需要新增一列overall_cor，存放每个id组内values_1与values_2的相关系数。

计算组内均值时，以下两种方法都能正常工作：

sample_df['overall_mean'] = sample_df.groupby('id')['values_1'].transform('mean')
sample_df['overall_mean'] = sample_df.groupby('id')['values_1'].mean()

但尝试用类似方式计算相关系数时，两种写法均报错：

sample_df['overall_cor'] = sample_df.groupby('id')['values_1','values_2'].transform('corr')
sample_df['overall_cor'] = sample_df.groupby('id')['values_1','values_2'].corr()

请问如何解决？为什么适用于均值的方法不能直接用于相关系数？

1. 聚合结果的形状差异 #

• 计算均值时，groupby('id')['values_1'].mean()对每个组的单列数据计算均值，返回的是每个组对应一个标量的Series，索引与id匹配，可通过广播或索引对齐直接赋值给原数据的新列。
• 而groupby('id')['values_1','values_2'].corr()对每个组的两列数据计算相关系数，返回的是每个组对应一个2×2的相关系数矩阵，整体为多层索引的DataFrame，形状与原数据完全不匹配，无法直接赋值。

2. transform的要求限制 #

transform方法要求传入的函数对每个组处理后，返回结果必须与组内行数长度一致（可以是标量，会自动广播到组内所有行；也可以是等长的Series/数组）。但corr函数返回的是矩阵，不符合transform的输入输出规则，因此直接调用transform('corr')会报错。

核心思路是先计算每个组的相关系数标量，再映射回原数据的每一行，以下是三种常用实现方式：

方法1：groupby.apply计算后合并 #

通过apply对每个组计算values_1和values_2的皮尔逊相关系数，得到id对应标量值后，用merge合并到原数据：

# 计算每个id的相关系数
corr_result = sample_df.groupby('id').apply(
    lambda group: group['values_1'].corr(group['values_2'])
).reset_index(name='overall_cor')

# 合并回原数据
sample_df = sample_df.merge(corr_result, on='id', how='left')

方法2：map方法直接映射 #

先构建id到相关系数的映射Series，再用map将原数据的id列映射为对应系数：

# 构建id与相关系数的映射
corr_map = sample_df.groupby('id').agg(
    overall_cor=lambda x: x['values_1'].corr(x['values_2'])
)

# 映射到原数据
sample_df['overall_cor'] = sample_df['id'].map(corr_map['overall_cor'])

方法3：transform结合自定义函数 #

如果一定要用transform，可自定义函数将相关系数标量广播为与组内行数一致的数组：

sample_df['overall_cor'] = sample_df.groupby('id').transform(
    lambda group: [group['values_1'].corr(group['values_2'])] * len(group)
)['values_1']

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