如何在DataFrame中计算两列相关系数并添加为新列（含分组场景）

问题1：为整个DataFrame的两列添加全局相关系数作为新列 #

如果你需要的是整个DataFrame中两列的全局相关系数（也就是一个单一值，广播到每一行），直接计算后赋值即可，pandas会自动把标量值扩展到整个列：

import pandas as pd

# 示例DataFrame
df = pd.DataFrame({
    'A': [1, 2, 3, 4, 5],
    'B': [2, 4, 5, 7, 9]
})

# 计算A和B的相关系数，作为新列添加
df['Corr_Coeff'] = df['A'].corr(df['B'])

这样Corr_Coeff列的每一行都会是A、B两列整体的相关系数值。

问题2：按ID分组计算每组的A、B相关系数，并映射回原DataFrame #

针对你给出的分组需求，核心是先按ID计算每组的相关系数，再把这个值关联到原DataFrame的每一行。这里有两种简洁的实现方式：

方法1：先分组计算，再用map关联

import pandas as pd

# 你的原始DataFrame
data = {
    'ID': [1, 1, 2, 2, 2],
    'A': [5, 3, 4, 7, 9],
    'B': [7, 4, 5, 6, 1]
}
df = pd.DataFrame(data)

# 第一步：按ID分组，计算每组A和B的相关系数
corr_per_id = df.groupby('ID').apply(lambda group: group['A'].corr(group['B']))

# 第二步：用map把每个ID对应的相关系数映射回原DataFrame
df['Corr_Coeff'] = df['ID'].map(corr_per_id)

方法2：用transform一步到位

transform方法会自动将分组计算的结果扩展为与原DataFrame行数一致的格式，直接赋值即可：

df['Corr_Coeff'] = df.groupby('ID')[['A', 'B']].transform(lambda x: x['A'].corr(x['B']))

运行后你会得到期望的结果：

ID  A  B  Corr_Coeff
0   1  5  7    0.999999
1   1  3  4    0.999999
2   2  4  5   -0.970725
3   2  7  6   -0.970725
4   2  9  1   -0.970725

简单解释下：两种方法都是先计算每个ID组内A和B的相关系数，再把这个值匹配到该ID的所有行上，区别只是实现方式的简洁度不同。

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