如何高效向量化计算两个3D矩阵的逐元素欧氏距离？

嘿，这个问题用numpy的广播机制就能完美解决，比写嵌套循环高效太多了！我来给你一步步拆解：

向量化核心实现 #

numpy的广播机制是处理这种维度扩展计算的利器，不用手动写循环就能让不同维度的矩阵自动匹配计算。直接看代码：

import numpy as np

# 先模拟输入矩阵（你可以替换成自己的真实数据）
matrix1 = np.random.rand(15, 3)  # shape (15, 3)
matrix2 = np.random.rand(23, 3)  # shape (23, 3)

# 利用np.newaxis扩展维度，触发广播
matrix_distance = (matrix1[:, np.newaxis, :] - matrix2[np.newaxis, :, :]) ** 2

代码解释： #

• matrix1[:, np.newaxis, :] 把原来的(15,3)矩阵变成**(15, 1, 3)**，增加了一个中间维度
• matrix2[np.newaxis, :, :] 把原来的(23,3)矩阵变成**(1, 23, 3)**，增加了第一个维度
• 当两个矩阵做减法时，numpy会自动把维度为1的部分扩展成对应对方的维度：(15,1,3) → (15,23,3)，(1,23,3) → (15,23,3)
• 最后逐元素平方，得到的matrix_distance就是你需要的**(15,23,3)**形状的矩阵，其中matrix_distance[:,:,0]就是第一通道所有元素的平方差，完全符合你的示例要求！

如果你觉得np.newaxis不够直观，也可以用np.expand_dims来写，效果完全一样：

matrix1_expanded = np.expand_dims(matrix1, axis=1)  # 等价于[:, np.newaxis, :]
matrix2_expanded = np.expand_dims(matrix2, axis=0)  # 等价于[np.newaxis, :, :]
matrix_distance = (matrix1_expanded - matrix2_expanded) ** 2

非向量化的循环写法（对比参考） #

你提到的非向量化解法雏形大概是这样的，用嵌套循环逐个计算：

# 初始化结果矩阵
matrix_distance_loop = np.zeros((15, 23, 3))

# 嵌套循环遍历所有元素对
for i in range(len(matrix1)):
    for j in range(len(matrix2)):
        matrix_distance_loop[i, j] = (matrix1[i] - matrix2[j]) ** 2

为什么向量化更好？ #

• 速度快：numpy的广播操作是底层用C实现的，比Python层面的for循环快几个数量级，尤其是当矩阵规模变大时（比如1000×3和2000×3的矩阵），差距会非常明显
• 代码简洁：一行代码搞定，可读性更高，也不容易出错
• 内存高效：广播不需要额外复制大量数据，而是在计算时动态扩展维度，内存占用比手动循环更低

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