Блог

Допустим, у нас есть 4D-массив A формы (n, m, g, h) и 1D-массив B формы (n).

Итак, я хочу выполнить операцию power для каждого 3D подмассива A (m, g, h) с каждым элементом B (который является int).

 A = np.arange(24).reshape(3, 2, 2, 2)
>>>array([[[[ 0,  1],
         [ 2,  3]],

        [[ 4,  5],
         [ 6,  7]]],


       [[[ 8,  9],
         [10, 11]],

        [[12, 13],
         [14, 15]]],


       [[[16, 17],
         [18, 19]],

        [[20, 21],
         [22, 23]]]]
B = np.arange(3)
>>>array([0, 1, 2])
  

Результат, который я хочу, это:

 C = somefunc(A, B) # just an example, can be anything
>>>array([[[[ (ignore, 0^0),  1],
         [ 1,  1]],

        [[ 1,  1],
         [ 1,  1]]],


       [[[ 8,  9],
         [10, 11]],

        [[12, 13],
         [14, 15]]],


       [[[256, 289],
         [324, 361]],

        [[400, 441],
         [484, 529]]]]
  

Каков наилучший способ получить желаемый результат?

Я думал о чем-то вроде:

 A = [np.power(A[i, :, :, :], B[i]) for i in range(B.size)]
  

Но это было бы довольно неэффективно, если B.size — большое число. Есть еще идеи?

Это сработает:

 C = A ** B[:,None,None,None]
print (C)
  

Вывод:

 [[[[  1   1]
   [  1   1]]

  [[  1   1]
   [  1   1]]]


 [[[  8   9]
   [ 10  11]]

  [[ 12  13]
   [ 14  15]]]


 [[[256 289]
   [324 361]]

  [[400 441]
   [484 529]]]]