Question

我使用 python 2.7 和 numpy 1.9 。
我有3种方法将转换应用于几个numpy数组。

def sum_arrays2(a, b):
    c = np.zeros(a.shape)
    c[:, 0:-1] = (a[:, 1:] + b[:, 0:-1]) ** 2
    return a[0:-1, 1:] + c[1:, 0:-1]

def sum_arrays3(a, b):
    c = np.zeros(a.shape)
    c[:, :, 0:-1] = (a[:, :, 1:] + b[:, :, 0:-1]) ** 2
    return a[0:-1, :, 1:] + c[1:, :, 0:-1]

def sum_arrays4(a, b):
    c = np.zeros(a.shape)
    c[:, :, :, 0:-1] = (a[:, :, :, 1:] + b[:, :, :, 0:-1]) ** 2
    return a[0:-1, :, :, 1:] + c[1:, :, :, 0:-1]

如你所见，它们非常相似。唯一的区别是所需的输入数组大小根据我的数据大小，我必须调用第一个，第二个或第三个。

实际上我必须做这样的事情：

if a.ndims == 2:
  result = sum_arrays2(a, b)
elif a.ndims == 3:
  result = sum_arrays3(a, b)
elif a.ndims == 4:
  result = sum_arrays4(a, b)

如何制作更通用的方法来计算n维输入？

我发现的唯一解决方案是：

def n_size_sum_arrays(a, b):
    c = np.zeros(a.shape)
    c[(Ellipsis, np.r_[0:c.shape[-1]-1])] = (a[(Ellipsis, np.r_[0:a.shape[-1]])] + b[(Ellipsis, np.r_[0:b.shape[-1]])]) ** 2
    return a[(r_[0:a.shape[0]-1], Ellipsis, np.r_[1:a.shape[-1]])] + c[(np.r_[1:c.shape[0]], Ellipsis, np.r_[0:c.shape[-1]-1])]

但它绝对不清楚，我不确定它是否正确。

有更好的方法吗？

Answer 1

您可以执行以下操作：

def sum_arrays(a, b):
    c = np.zeros(a.shape)
    c[..., :-1] = (a[..., 1:] + b[..., :-1]) ** 2
    return a[:-1, ..., 1:] + c[1:, ..., :-1]

编写一个访问参数化维度numpy数组的方法

1 个答案: