如何高效计算两个DataFrame列间的距离矩阵？

原始数据 #

我有两个DataFrame，每一列对应不同人员的数值数据，具体定义如下：

import numpy as np
import pandas as pd
import math

df1 = pd.DataFrame({
    "Anna":[1.5,-2,2.5],
    "Bob":[2.5,-3,3.5],
    "Cam":[3.5,-4,4.5]})
df2 = pd.DataFrame({
    "Dave":[1,-2.5,2],
    "Emma":[2,-3.5,3],
    "Fred":[3,-4.5,4]})

print(df1)
#    Anna  Bob  Cam
# 0   1.5  2.5  3.5
# 1  -2.0 -3.0 -4.0
# 2   2.5  3.5  4.5

print(df2)
#    Dave  Emma  Fred
# 0   1.0   2.0   3.0
# 1  -2.5  -3.5  -4.5
# 2   2.0   3.0   4.0

原有实现（双重循环） #

当前通过双重循环计算df1中每个人员与df2中每个人员的欧氏距离矩阵：

results = []

for n1 in df1.columns:
    results.append([])
    for n2 in df2.columns:
        results[-1].append(math.dist(df1[n1], df2[n2]))

res_df = pd.DataFrame(results)
res_df.columns = df1.columns
res_df.index = df2.columns

print(res_df)
#         Anna        Bob         Cam
# Dave  0.866025    1.658312    3.278719
# Emma  2.179449    0.866025    1.658312
# Fred  3.840573    2.179449    0.866025

更高效的实现方式 #

Python显式循环的开销较大，数据量增长时效率会明显下降，以下两种向量化/库函数实现能大幅提升计算效率：

方法1：NumPy广播向量化计算 #

利用NumPy的广播机制，一次性完成所有向量对的距离计算，彻底规避Python循环开销：

# 将DataFrame转置，使每行对应一个人员的数值向量
arr1 = df1.T.values  # shape: (3, 3)，对应df1的3个人员
arr2 = df2.T.values  # shape: (3, 3)，对应df2的3个人员

# 通过广播计算所有向量对的欧氏距离
# arr1[:, np.newaxis] 将维度扩展为(3,1,3)，与arr2广播后得到(3,3,3)的差矩阵
# 对最后一维求和（计算平方和），再开平方得到欧氏距离
distances = np.sqrt(np.sum((arr1[:, np.newaxis] - arr2) ** 2, axis=2))

# 转换为符合要求的DataFrame
res_df = pd.DataFrame(distances.T, index=df2.columns, columns=df1.columns)
print(res_df)

方法2：SciPy的cdist函数 #

scipy.spatial.distance.cdist是专门用于计算两个向量集合间距离的优化函数，底层由C实现，是大样本量下效率最优的方案：

from scipy.spatial.distance import cdist

arr1 = df1.T.values
arr2 = df2.T.values

# cdist(X, Y) 计算X中每个向量到Y中每个向量的距离，结果shape为(len(X), len(Y))
distances = cdist(arr2, arr1)

res_df = pd.DataFrame(distances, index=df2.columns, columns=df1.columns)
print(res_df)

效率对比说明 #

• 双重循环：时间复杂度为O(M*N*D)（M、N为两个DataFrame的列数，D为行数），且Python循环本身有额外开销，数据量越大效率越低。
• 向量化/SciPy方法：将计算逻辑转移到底层优化的C层执行，避免Python循环开销，当列数超过100时，效率提升可达数十倍甚至上百倍。

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