使用 numpy 向量化此操作

Question

为了介绍这个问题，假设我有一个存储 N 个点的 2D 数组（Nx3）。另一方面，我有一个 1D 数组 (M, 3)，其中包含单个 3D 参考点的坐标，因此 M=1

对于这种情况，如果我想计算一个新的一维数组（N）以及到参考点的所有距离，我只需要这样做：

import numpy as np

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

print(np.linalg.norm(points - references, axis=1))

这有效。但是，当我有超过 1 个参考点 (

M>1

) 时，numpy 无法广播该操作。现在我只是做与前面描述的相同的操作，但是使用 for 循环在参考点上进行逐一迭代，这是相当低效的。

当有

M > 1

时，如何以矢量化方式实现上述行为？我之前考虑过的事情：

• 使用

  numba.njit

  去除python解释的效果

  for

  。但我想不做

  numba

• 比如将点数组广播为 3D (MxNx3) 并将参考数组广播为 3D (NxMx3)，以便我可以操作它们。但是，如果我没记错的话，我会将点数组的大小乘以 M，并将参考数组的大小乘以 N。如果这是解决方案，我也不知道该怎么做。

Answer 1

简单地说：

np.linalg.norm(references - points[:, None], axis=-1)

这为您提供了所有距离的

M,N

数组。

示例：

np.random.seed(0)  # reproducible setup
N = 4
M = 2
points = np.random.random((N, 3))
references = np.random.random((M, 3))

>>> np.linalg.norm(references - points[:, None], axis=-1)
array([[0.57216693, 0.86543108],
       [0.76350759, 0.63809564],
       [0.90274337, 0.94847268],
       [0.51150232, 0.88049546]])

注意：如果您正在处理大量点，并且不一定对完整的距离列表感兴趣，而是对最近邻的子集感兴趣，那么您应该考虑使用

scipy.spatial.KDTree

。这会更快并且内存效率更高。

Answer 2

当您获取数组的差异时进行广播，然后将

np.linalg.norm

应用于最后一个轴（我明确写了 2，但您也可以使用 -1）。

np.linalg.norm(points[:,None] - references[None,:], axis=2)

请注意，当您的版本为您提供一维向量时，当 M = 1 时，您最终会得到一个列向量。如果您想要一个行向量，请交换

points

和

references

索引。这也将产生转置最终结果的效果。我不确定你喜欢哪个方向。