如何在没有for循环的情况下计算numpy数组

Question

我有两个形状分别为A = (226, 250)和B = (195, 195, 250)的numpy数组。我想计算一个新数组C：

for i in range(A.shape[0]):
    C = A[i, :].reshape(1, 1, -1) - B        

是否有另一种无需迭代过程即可计算C的方法？

Answer 1

似乎想要

C = A[:, None, None, :] - B

None是np.newaxis的别名。它将在您将其放置在索引中的位置插入新轴。就像

C = A.reshape(A.shape[0], 1, 1, A.shape[1])

Answer 2

您可以使用apply_along_axis

C = np.apply_along_axis(lambda x: x.reshape(1, 1, -1) - B, 1, A)

但是请记住，整个数据集都是在ram中处理的。这可能会导致RAM问题。

编辑： 使用C = A[:, None, None, :] - B的解决方案在RAM上存在相同的问题。

另一种方法是保持循环并使用numba装饰器。这样可以将循环速度提高到接近c的速度，并避免淹没您的RAM。