Python如何提高numpy阵列性能？

我有一个全局numpy.array数据，它是一个200*200*3的3d数组，在3d空间中包含40000个点

我的目标是计算从每个点到单位立方体四个角的距离（（0,0,0），（1,0,0），（0,1,0），（0,0,1）），以便确定哪个角离它最近

def dist(*point):
    return np.linalg.norm(data - np.array(rgb), axis=2)

buffer = np.stack([dist(0, 0, 0), dist(1, 0, 0), dist(0, 1, 0), dist(0, 0, 1)]).argmin(axis=0)

我编写了这段代码并进行了测试，每次运行大约需要10毫秒。 我的问题是如何提高这段代码的性能，在不到1ms的时间内更好地运行。

您可以使用-

运行时测试

# Original approach
def org_app():
    return np.stack([dist(0, 0, 0), dist(1, 0, 0), \
       dist(0, 1, 0), dist(0, 0, 1)]).argmin(axis=0)

时间安排-

In [170]: data = np.random.rand(200,200,3)

In [171]: %timeit org_app()
100 loops, best of 3: 4.24 ms per loop

In [172]: %timeit cdist(data.reshape(-1,3), uc).argmin(1).reshape(data.shape[0],-1)
1000 loops, best of 3: 1.25 ms per loop

---

一个立方体不是有四个角吗？@JohnZwinck只需要计算到其中四个角的距离。

In [170]: data = np.random.rand(200,200,3)

In [171]: %timeit org_app()
100 loops, best of 3: 4.24 ms per loop

In [172]: %timeit cdist(data.reshape(-1,3), uc).argmin(1).reshape(data.shape[0],-1)
1000 loops, best of 3: 1.25 ms per loop