基于Pandas两列列表值生成皮尔逊相关系数列的实现优化

1. Lambda 直接实现版本 #

你需要的lambda调用写法可以直接实现相同效果，代码如下：

from scipy import stats
import pandas as pd

# 一行实现，直接取pearsonr返回的第一个值（相关系数）
df['Correlation'] = df.apply(lambda row: stats.pearsonr(row['A'], row['B'])[0], axis=1)

该写法和你自定义函数的逻辑完全等价，计算结果一致。

2. 效率优化建议 #

apply(axis=1)效率低是pandas的共性问题，逐行处理会产生大量行对象的额外开销，数据量越大性能损耗越明显，可以根据你的数据规模选择下面的优化方案：

方案一：列表推导替代apply（中小数据量） #

提前把两列的列表转成numpy数组，再用列表推导逐行计算，比apply性能高30%~50%：

import numpy as np

a_data = np.array(df['A'].tolist())
b_data = np.array(df['B'].tolist())
df['Correlation'] = [stats.pearsonr(a, b)[0] for a, b in zip(a_data, b_data)]

方案二：纯numpy向量化计算（大数据量最优） #

完全避免Python层面的循环，基于皮尔逊相关系数的数学公式用numpy矩阵操作实现，性能是逐行计算的10倍以上：

a_arr = np.array(df['A'].tolist())
b_arr = np.array(df['B'].tolist())

# 按行计算均值
a_mean = a_arr.mean(axis=1, keepdims=True)
b_mean = b_arr.mean(axis=1, keepdims=True)

# 去均值处理
a_centered = a_arr - a_mean
b_centered = b_arr - b_mean

# 计算协方差、标准差乘积
cov = (a_centered * b_centered).sum(axis=1)
std_product = np.sqrt((a_centered ** 2).sum(axis=1) * (b_centered ** 2).sum(axis=1))

# 生成结果，处理除0异常（某行数值全相同的场景）
df['Correlation'] = np.where(std_product == 0, np.nan, cov / std_product)

注：如果你的数据中不存在单行列所有值都相同的场景，可以去掉异常处理逻辑直接计算

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