python - Numpy，避免在 3d 数组差异嵌套求和中循环

Question

我有一个简单的 Numpy 问题:我有 3d 坐标，我想用以下函数计算两个不同配置之间的重叠

def Overlap(rt, r0,a):
    s=0
    for i in range(len(rt)):
        s+=(( pl.norm(r0[i]-rt ,axis=1) <=a).astype('int')).sum()
    return s`

其中 rt 和 r0 代表两个 m x 3 表，配置。实际上，它计算第一个配置中的向量与第二个配置中的任何其他向量之间的距离，检查阈值 a，并在遍历所有位置后返回总和。有没有避免显式 for 循环的聪明方法？我有一种感觉，复杂性无法真正改变，但也许有一种方法可以避免 native for construct 的缓慢。

Answer 1

以下情况如何:

from scipy.spatial.distance import cdist
import numpy as np
overlap = np.sum(cdist(rt, r0) <= a)

当 m 在我的机器上为 1000 时，速度大约快 9 倍。对于小型阵列要快得多