如何在没有硬编码维数的情况下枚举ND阵列的所有元素？

Question

如果我知道尺寸的数量，比如3，我会用3个嵌套循环对其进行硬编码：

for i in range(A.shape[0]):
   for j in range(A.shape[1]):
      for k in range(A.shape[2]):
          A[i,j,k] = some_formula(i, j, k)

但如果我不知道尺寸的数量怎么办？我是否仍然可以在知道每个循环上的所有索引的情况下枚举数组？

Answer 1

如果您的功能广播，您可以使用numpy.fromfunction：

B = numpy.fromfunction(some_formula, A.shape, dtype=int)

如果你的功能没有广播，你可以使用numpy.vectorize和numpy.fromfunction，但它的效率会比你的功能自然广播的效率低很多：

B = numpy.fromfunction(numpy.vectorize(some_formula), A.shape, dtype=int)