如何向量化计算两个矩阵逐行的np.linalg.norm以优化循环性能

优化后完整代码 #

import numpy as np

# 原示例数组定义不变
arr1 = np.random.random(300).reshape((-1,3))
arr2 = np.random.random(450).reshape((-1,3))

# 向量化计算，完全替代双重循环
dists = np.linalg.norm(arr1[:, np.newaxis] - arr2, axis=2, ord=2)
result_arr = 1 / (1 + dists ** (1/3))

实现逻辑说明 #

• 利用Numpy广播机制调整数组维度：arr1[:, np.newaxis]会将形状为(N, 3)的arr1扩展为(N, 1, 3)，和形状为(M, 3)的arr2做减法时，自动广播生成形状为(N, M, 3)的差值数组，其中[i,j,:]位置的元素就是原循环中arr1[i] - arr2[j]的计算结果
• 调用np.linalg.norm时指定axis=2，直接在坐标维度上逐对计算L2范数，生成形状为(N, M)的距离矩阵，和原循环逐元素计算的结果完全一致
• 最后对整个距离矩阵做逐元素算术运算，完全兼容原有的运算逻辑，所有计算都在Numpy底层C实现中完成，没有Python层面的循环开销，性能相比原双重循环提升可达百倍以上。

可选内存优化方案 #

如果处理的数组规模极大，广播生成(N, M, 3)的中间数组会占用过多内存，可以改用scipy的cdist方法直接计算距离矩阵，无需生成中间差值数组：

from scipy.spatial.distance import cdist
dists = cdist(arr1, arr2, metric='euclidean')
result_arr = 1 / (1 + dists ** (1/3))

该方法计算结果和前面的向量化方案完全一致，内存占用更低，运算速度更快。

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