NDarray 中的 Numpy 一维数组操作

Question

假设我们有一个形状为 (n, m, g, h) 的四维数组 A 和一个形状为 (n) 的一维数组 B。

所以我想对 A (m, g, h) 的每个 3D 子数组和 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 是一个很大的数字，那将是非常低效的。还有其他想法吗？

Answer 1

这将起作用：

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]]]]