如何用矩阵运算代替for循环？

Question

for x in range(0, width):
    for y in range(0, width):
        D[x,y] = dist(A[x,y], B[x,y])

def dist(a, b):
    return ((a-b)**2).sum()

A和B中的值是形状为（3,3）的矩阵

如何使用矩阵运算而不是for循环获得相同的结果？ 以下代码不起作用。

D[[range(0,height)],[range(0,width)]] = dist(A[[range(0,height)],[range(0,width)], 
                                             B[[range(0,height)],[range(0,width)])]

Answer 1

我对此的解释是A和B都是尺寸为h * w * 3 * 3的NumPy数组。

使用数组运算查找每个A[i,j,:,:]的3×3矩阵B[i,j,:,:]和i,j之间的平方距离的最简单方法是：

distances = np.sum((A-B)**2, axis=(2,3))

即只是对所有东西进行平方并在最后两个轴上求和。

这涉及到创建一些不必要的临时数组（A-B，然后是其正方形），因此，如果您需要 extreme 效率，则可以使用另一个选项，例如Numba jitted for循环，为np.vdot和A[i,j,:,:]的每一对运行B[i,j,:,:]。但这不太可能是您的瓶颈。