如何在不对维数进行硬编码的情况下枚举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)