Pandas中行子集间距离计算的高效实现方案问询

有如下格式的Pandas DataFrame：

S1       S2      id    var
  0      1.2      3.2     A1    A
  1      3.4      0.4     A2    A
  2      -2.3     1.2     A3    A
  3      0.1      -1.3    B1    B
  4      4.5      1.3     B2    B
  5      -2.3     -1.2    C1    C

需要计算不同组间所有行的两两欧氏距离平均值：

• dist_AB = ((A1-B1距离) + (A1-B2距离) + (A2-B1距离) + (A2-B2距离))/4
• dist_AC = ((A1-C1距离) + (A2-C1距离) + (A3-C1距离))/3
• dist_BC = ((B1-C1距离) + (B2-C1距离))/2

原实现采用多层循环逐行计算，处理数万行数据时效率极低，需更高效的实现方式。

核心思路是用向量化运算替代逐行循环，利用数值计算库的批量处理能力，大幅提升运行效率。

方法1：使用scipy.spatial.distance.cdist批量计算 #

cdist是专门用于计算两组样本两两距离的函数，底层经过优化，代码简洁且性能优异：

import io
import numpy as np
import pandas as pd
from scipy.spatial.distance import cdist
import itertools

TESTDATA="""
S1       S2      id    var
1.2      3.2     A1    A
3.4      0.4     A2    A
-2.3     1.2     A3    A
0.1      -1.3    B1    B
4.5      1.3     B2    B
-2.3     -1.2    C1    C
"""
df = pd.read_csv(io.StringIO(TESTDATA), sep="\s+")

# 按var分组提取特征列数据
grouped_features = df.groupby('var')[['S1', 'S2']]
group_names = list(grouped_features.groups.keys())

# 遍历所有组对，计算平均距离
result_list = []
for v1, v2 in itertools.combinations(group_names, 2):
    # 获取两组的特征矩阵
    group1_data = grouped_features.get_group(v1).values
    group2_data = grouped_features.get_group(v2).values
    # 计算两两欧氏距离矩阵
    distance_matrix = cdist(group1_data, group2_data, metric='euclidean')
    # 计算该组对的平均距离
    avg_distance = distance_matrix.mean()
    result_list.append({'var': v1+v2, 'dist': avg_distance})

# 转换为DataFrame输出
distances_df = pd.DataFrame(result_list)
print(distances_df)

输出结果与原代码完全一致：

var      dist
0   AB  3.973345
1   AC  4.647527
2   BC  4.823540

方法2：纯Numpy广播实现（无SciPy依赖） #

如果不想引入SciPy依赖，可利用Numpy的广播机制实现批量距离计算：

import io
import numpy as np
import pandas as pd
import itertools

TESTDATA="""
S1       S2      id    var
1.2      3.2     A1    A
3.4      0.4     A2    A
-2.3     1.2     A3    A
0.1      -1.3    B1    B
4.5      1.3     B2    B
-2.3     -1.2    C1    C
"""
df = pd.read_csv(io.StringIO(TESTDATA), sep="\s+")

grouped_features = df.groupby('var')[['S1', 'S2']]
group_names = list(grouped_features.groups.keys())

result_list = []
for v1, v2 in itertools.combinations(group_names, 2):
    group1_data = grouped_features.get_group(v1).values
    group2_data = grouped_features.get_group(v2).values
    # 广播计算两组样本的差值
    diff = group1_data[:, np.newaxis, :] - group2_data[np.newaxis, :, :]
    # 计算欧氏距离矩阵：差值平方和开根号
    distance_matrix = np.sqrt(np.sum(diff ** 2, axis=2))
    avg_distance = distance_matrix.mean()
    result_list.append({'var': v1+v2, 'dist': avg_distance})

distances_df = pd.DataFrame(result_list)
print(distances_df)

效率说明 #

• 原循环实现：逐行迭代的Python级循环开销极大，时间复杂度为O(M*N)（M、N为两组样本量），处理大规模数据时性能断崖式下降。
• 向量化实现：利用C语言优化的底层数值计算，同样的时间复杂度但实际运行速度比循环快100~1000倍，可轻松处理数万行甚至更大规模的数据。

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