如何使用迭代函数为同一维度的numpy数组分配新的不同值

Posted 2021-04-08

我有以下多维numpy数组：

a = np.arange(16).reshape(2,2,2,2)

我想为某个维度的每个元素分配新的不同数组值，例如第四维我使用以下代码：

    for i in range(a.shape[3]):
          if i == 0:
               for t in np.nditer(a[:,:,:,i], op_flags = ['readwrite']):
                    t[...] = t*2
          if i == 1:
               for t in np.nditer(a[:,:,:,i], op_flags = ['readwrite']):
                    t[...] = t*3

print(a)
print(a.shape)

输出显示为

[[[[ 0  1]
   [ 4  3]]

   [[ 8  5]
   [12  7]]]


 [[[16  9]
   [20 11]]

  [[24 13]
   [28 15]]]]
[[[[ 0  3]
   [ 4  9]]

  [[ 8 15]
   [12 21]]]


 [[[16 27]
   [20 33]]

  [[24 39]
   [28 45]]]]
(2, 2, 2, 2)
(2, 2, 2, 2)

据我所知，它迭​​代数组并在第一个i它分配新值然后在下一个i它分配新值并创建新数组，除了第一个i的第一个数组之外的两个新值，这就是为什么我有两个数组的原因一个变量。我只关心最后一个数组，其中所有值都分配了新值。我怎么才能提取最后一个数组。或者还有另一个代码，这个任务更简单，更省时？

答案

你可以这样做：

>>> a[:,:,:,0] = a[:,:,:,0]*2
>>> a[:,:,:,1] = a[:,:,:,1]*3
>>> a
array([[[[ 0,  3],
         [ 4,  9]],

        [[ 8, 15],
         [12, 21]]],


       [[[16, 27],
         [20, 33]],

        [[24, 39],
         [28, 45]]]])

甚至这个：

>>> a[:,:,:,0] *= 2
>>> a[:,:,:,1] *= 3