如何用Numpy向量化实现多参考点与点集的距离计算？

问题背景 #

现有N×3的二维数组存储N个3D点，M×3的数组存储M个3D参考点。当M=1时，可通过Numpy广播结合np.linalg.norm直接计算所有点到参考点的距离，但M>1时直接广播失效，循环遍历参考点效率极低，需要不依赖numba.njit的向量化实现方案。

向量化实现方案 #

方法1：扩维广播+范数计算 #

核心是通过扩维让两个数组满足广播规则，生成(N,M,3)的差值数组后，对最后一维求范数得到距离矩阵：

import numpy as np

N = 100
M = 5
points = np.random.random((N, 3))
references = np.random.random((M, 3))

# 给points添加中间维度，使其形状变为(N,1,3)，与(M,3)的references广播
distances = np.linalg.norm(points[:, np.newaxis, :] - references, axis=2)
# 也可以用np.expand_dims替代np.newaxis，可读性更强：
# distances = np.linalg.norm(np.expand_dims(points, axis=1) - references, axis=2)

print(distances.shape)  # 输出 (100,5)，对应每个点到每个参考点的距离

方法2：矩阵运算推导欧氏距离 #

利用欧氏距离的平方公式||a-b||² = ||a||² + ||b||² - 2a·b，通过矩阵乘法实现，避免创建大尺寸中间数组，内存效率更高：

import numpy as np

N = 100
M = 5
points = np.random.random((N, 3))
references = np.random.random((M, 3))

# 计算每个点的模长平方，保持维度便于广播
points_norm = np.sum(points**2, axis=1, keepdims=True)  # 形状(N,1)
# 计算每个参考点的模长平方
refs_norm = np.sum(references**2, axis=1, keepdims=True)  # 形状(M,1)
# 计算点与参考点的点积矩阵
dot_product = points @ references.T  # 形状(N,M)
# 推导距离并开根号
distances = np.sqrt(points_norm + refs_norm.T - 2 * dot_product)

print(distances.shape)  # 输出 (100,5)

方案对比 #

• 方法1代码简洁直观，但N、M较大时会生成(N,M,3)的中间数组，内存占用较高。
• 方法2通过矩阵运算减少内存开销，计算效率更优，适合大规模数据场景。

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