如何用NumPy向量化计算两个三维坐标矩阵的范数差矩阵？

问题背景 #

给定形状为(n, 3)的三维坐标矩阵v和形状为(m, 3)的三维坐标矩阵u，需要生成形状为(n, m)的矩阵M，其中M[i, j]是v[i]与u[j]之间的欧氏距离（第二向量范数）。原嵌套循环实现虽然直观，但处理大数据集时效率极低，以下是两种高效的NumPy向量化实现方案。

原循环实现（供参考） #

import numpy as np

n, m = 100, 200
M = np.zeros((n, m))
v = np.random.rand(n, 3)
u = np.random.rand(m, 3)

for i in range(n):
    for j in range(m):
        M[i, j] = np.linalg.norm(v[i] - u[j], ord=2)

方案一：广播+直接范数计算 #

利用NumPy的广播机制，将两个矩阵扩展维度后直接计算差的范数，代码简洁直观：

import numpy as np

n, m = 100, 200
v = np.random.rand(n, 3)
u = np.random.rand(m, 3)

# 扩展维度后广播相减，再对每个3维向量计算范数
M = np.linalg.norm(v[:, np.newaxis, :] - u[np.newaxis, :, :], axis=2)

原理：

• v[:, np.newaxis, :]将v从(n, 3)转换为(n, 1, 3)；
• u[np.newaxis, :, :]将u从(m, 3)转换为(1, m, 3)；
• 两者相减时触发广播，得到形状为(n, m, 3)的差矩阵；
• 对最后一个轴（axis=2）计算L2范数，直接得到(n, m)的距离矩阵。

方案二：欧氏距离公式展开（内存更高效） #

通过数学公式展开计算，避免创建大尺寸的中间差矩阵，适合处理超大数据集：

欧氏距离的平方满足：
$$||v_i - u_j||^2 = ||v_i||^2 + ||u_j||^2 - 2v_i \cdot u_j$$

基于此的实现代码：

import numpy as np

n, m = 100, 200
v = np.random.rand(n, 3)
u = np.random.rand(m, 3)

# 计算每个点的L2范数平方，并调整维度以支持广播
v_norm_sq = np.sum(v ** 2, axis=1)[:, np.newaxis]  # 形状(n, 1)
u_norm_sq = np.sum(u ** 2, axis=1)[np.newaxis, :]  # 形状(1, m)

# 计算v与u的点积矩阵
dot_product = np.dot(v, u.T)  # 形状(n, m)

# 计算距离，用np.maximum避免浮点数精度导致的极小负数
M = np.sqrt(np.maximum(v_norm_sq + u_norm_sq - 2 * dot_product, 0))

原理：

• 先分别计算所有点的范数平方，通过维度调整实现广播相加；
• 点积矩阵dot_product存储了每对v[i]和u[j]的内积；
• 代入公式计算距离平方后开根号，用np.maximum确保根号内非负，避免因浮点数误差引发警告。

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