如何通过broadcasting机制快速计算两个dataframe每行的自定义距离

实现原理 #

• 假设你有两个DataFrame：df1（m行n列）和df2（k行n列），只需要将两者转为numpy数组后分别扩展一个维度，得到形状为(m, 1, n)和(1, k, n)的三维数组，利用numpy的广播特性，两者运算时会自动补齐为(m, k, n)的数组，刚好对应df1每一行和df2每一行的每一列的差值计算。
• 不同列指定不同距离规则时，直接对三维数组的最后一维（列维度）按索引分别做对应距离的计算，最后对列维度聚合求和即可得到每对行的总自定义距离。

示例代码 #

import pandas as pd
import numpy as np

# 构造示例数据
df1 = pd.DataFrame({
    "col1": [1, 3, 5], # 第一列用欧氏距离
    "col2": [2, 4, 6]  # 第二列用曼哈顿距离
})
df2 = pd.DataFrame({
    "col1": [2, 4],
    "col2": [1, 3]
})

# 转为数组并扩展维度适配广播规则
arr1 = df1.values[:, np.newaxis, :]
arr2 = df2.values[np.newaxis, :, :]

# 按列分别计算对应距离分量
col_distance = np.zeros((arr1.shape[0], arr2.shape[1], arr1.shape[2]))
# 第一列：欧氏距离分量（差的平方，后续可根据需求开根号）
col_distance[..., 0] = np.square(arr1[..., 0] - arr2[..., 0])
# 第二列：曼哈顿距离分量（差的绝对值）
col_distance[..., 1] = np.abs(arr1[..., 1] - arr2[..., 1])

# 聚合得到最终距离矩阵，total_dist[i][j] 对应df1第i行和df2第j行的总距离
# 若第一列需要标准欧氏距离，可改为 np.sqrt(col_distance[...,0]) + col_distance[...,1]
total_dist = col_distance.sum(axis=-1)

方案优势 #

• 全程为numpy向量化运算，没有显式的Python层循环，性能远高于逐行遍历计算的方案，数据量越大性能优势越明显。
• 扩展性强，新增列只需要对应新增列索引的距离计算逻辑即可。

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