矢量化三个嵌套循环 - NumPy

时间:2016-07-19 14:32:49

标签: python numpy vectorization

我有两个数组x.dim = (N,4)y.dim = (M, M, 2)以及一个函数f(a, b),它分别将KL维向量作为参数。我想获得一个数组res.dim = (N, M, M),以便

for n in range(N):
  for i in range(M):
    for j in range(M):
      res[n, i, j] = f(x[n], y[i, j])

在这种情况下无法使用apply。在此先感谢您的帮助!

def f(a, b):
    return max(0, 1 - np.sum(np.square(np.divide(np.subtract(b, a[0:2]), a[2:4]))))

1 个答案:

答案 0 :(得分:1)

这是使用NumPy broadcasting和切片 -

处理所列函数的矢量化方法
# Slice out relevant cols from x
x_slice1 = x[:,None,None,:2]
x_slice2 = x[:,None,None,2:4]

# Perform operations using those slices to correspond to iterative operations
out = np.maximum(0,1-(np.divide(y-x_slice1,x_slice2)**2).sum(3))