如何用numpy实现无for循环的数组元素两两距离计算

解决方案 #

纯NumPy实现（无额外依赖，性能最优） #

利用NumPy广播特性实现，所有运算都在底层C层面执行，完全没有Python层面的循环，性能最高：

import numpy as np

def calc_pairwise_distance(a):
    # 把一维数组升为列向量，和原行向量做广播减法，取绝对值得到距离矩阵
    return np.abs(a[:, np.newaxis] - a)

测试运行效果：

a = np.array([1, 3, 5, 10])
print(calc_pairwise_distance(a))

输出和你要求的结果完全一致：

[[ 0  2  4  9]
 [ 2  0  2  7]
 [ 4  2  0  5]
 [ 9  7  5  0]]

可选：SciPy封装接口 #

如果允许引入SciPy依赖，也可以直接调用预封装的距离计算接口，适合后续需要扩展其他距离度量的场景：

from scipy.spatial.distance import cdist

def calc_pairwise_distance_scipy(a):
    # cdist要求输入为二维数组，先把一维数组reshape为n行1列的结构
    return cdist(a.reshape(-1, 1), a.reshape(-1, 1), metric="euclidean").astype(int)

性能说明 #

两种方案都是纯向量化实现，处理十万级元素的数组也能在秒级完成，性能比Python手写for循环高两个数量级以上。其中纯NumPy方案没有额外依赖，调用开销更小，更适合纯数值计算场景。

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