按索引分组计算指定列与DataFrame其余列的快速相关系数

问题背景 #

现有一个尺寸为(109049, 29184)的DataFrame，索引包含22种不同取值，数据结构示例如下：

df:
      Ford Honda GM
index           
Sedan   4   1   8
Sedan   5   2   7
Sedan   6   3   6
Sedan   7   4   5
SUV     8   5   7
SUV     1   6   6
SUV     2   7   5
SUV     3   8   4

需求是按索引分组，计算Ford列与其他所有列的Spearman相关系数，输出格式要求如下：

index   SUV      Sedan
        Ford     Ford
Ford    1.00     1.0
Honda   -0.58    1.0
GM      0.58     -1.0

此前使用代码df.groupby('index').corr(method = 'spearman').reset_index()计算，但因数据量过大，运行超10小时仍未完成，需寻求高效实现方法。

高效解决方案 #

核心思路 #

原生groupby.corr会计算所有列对的相关系数，产生大量冗余计算（比如我们不需要Honda与GM的相关系数）。而Spearman相关系数本质是变量秩次的Pearson相关系数，因此可以针对性优化：

• 先对每个分组内的所有列计算秩次
• 仅计算Ford列秩次与其他列秩次的Pearson相关系数，直接得到目标Spearman结果

代码实现 #

import pandas as pd

# 1. 对每个分组内的所有列计算秩次（Spearman的核心依赖）
ranked_df = df.groupby(level='index').rank(method='average')

# 2. 提取Ford列的秩次，作为相关计算的基准
ford_ranks = ranked_df['Ford']

# 3. 按分组计算其他列与Ford列的Pearson相关系数（等价于原数据的Spearman系数）
result = ranked_df.groupby(level='index').apply(
    lambda group: group.corrwith(ford_ranks.loc[group.index], method='pearson')
).unstack(0)

# 调整列层级，匹配需求的输出格式
result.columns = pd.MultiIndex.from_tuples([(col, 'Ford') for col in result.columns])
result.index.name = None

# 可选：保留两位小数对齐输出格式
result = result.round(2)

print(result)

性能提升原因 #

• 计算量从全量列对的O(n²)降到仅目标列对的O(n)（n为列数），大幅减少冗余运算
• 秩次计算是pandas优化后的向量化操作，效率远高于全量相关系数计算
• 针对22个分组+29184列的场景，该方法能将计算时间压缩到原方法的几十分之一

额外优化建议 #

• 若内存不足，可尝试分批次处理单个分组，或用dask.dataframe实现并行计算
• 确保使用最新版本的pandas，新版本对groupby、rank等操作有显著性能优化

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