将numpy.arange广播到多维数组中

Question

假设我有一个零的3维数组

Y = np.zeros((2,3,4))

我希望通过第二维向它广播一维数组

X = np.arange(3)

这样得到的多维数组是

Y = [[[ 0. 0. 0. 0.]
      [ 1. 1. 1. 1.]
      [ 2. 2. 2. 2.]]

     [[ 0. 0. 0. 0.]
      [ 1. 1. 1. 1.]
      [ 2. 2. 2. 2.]]]

在没有循环（甚至列表理解）的情况下，最Python高效的方法是什么？ 我希望能够扩大尺寸，并使用numba模块运行代码，该模块无法使用np.tile或np.repeat进行编译

Answer 1

您可以首先重塑X ，以便可以将其广播到Y.shape ，然后使用np.broadcast_to方法：

Y = np.zeros((2,3,4))  
np.broadcast_to(X[:,None], Y.shape)

#array([[[0, 0, 0, 0],
#        [1, 1, 1, 1],
#        [2, 2, 2, 2]],

#       [[0, 0, 0, 0],
#        [1, 1, 1, 1],
#        [2, 2, 2, 2]]])

---

另一个选择是，创建一个形状与Y相同的新空数组，然后为其分配X值：

new_X = np.empty_like(Y)    
new_X[:] = X[:,None]

new_X
#array([[[ 0.,  0.,  0.,  0.],
#        [ 1.,  1.,  1.,  1.],
#        [ 2.,  2.,  2.,  2.]],

#       [[ 0.,  0.,  0.,  0.],
#        [ 1.,  1.,  1.,  1.],
#        [ 2.,  2.,  2.,  2.]]])

Answer 2

如果您始终以零数组开头，则只需添加x广播

>>> y + x[None, :, None]
array([[[ 0.,  0.,  0.,  0.],
        [ 1.,  1.,  1.,  1.],
        [ 2.,  2.,  2.,  2.]],

       [[ 0.,  0.,  0.,  0.],
        [ 1.,  1.,  1.,  1.],
        [ 2.,  2.,  2.,  2.]]])

一般广播规则

要么

>>> np.add(y, x[None, :, None], out = y)
array([[[ 0.,  0.,  0.,  0.],
        [ 1.,  1.,  1.,  1.],
        [ 2.,  2.,  2.,  2.]],

       [[ 0.,  0.,  0.,  0.],
        [ 1.,  1.,  1.,  1.],
        [ 2.,  2.,  2.,  2.]]])
>>> y
array([[[ 0.,  0.,  0.,  0.],
        [ 1.,  1.,  1.,  1.],
        [ 2.,  2.,  2.,  2.]],

       [[ 0.,  0.,  0.,  0.],
        [ 1.,  1.,  1.,  1.],
        [ 2.,  2.,  2.,  2.]]])
>>> 

我不确定您是否可以使用分配进行广播，但可以做到：

>>> y[:] = x[None, :, None]
>>> y
array([[[ 0.,  0.,  0.,  0.],
        [ 1.,  1.,  1.,  1.],
        [ 2.,  2.,  2.,  2.]],

       [[ 0.,  0.,  0.,  0.],
        [ 1.,  1.,  1.,  1.],
        [ 2.,  2.,  2.,  2.]]])
>>> y[...] = x[None, :, None]
>>> y
array([[[ 0.,  0.,  0.,  0.],
        [ 1.,  1.,  1.,  1.],
        [ 2.,  2.,  2.,  2.]],

       [[ 0.,  0.,  0.,  0.],
        [ 1.,  1.,  1.,  1.],
        [ 2.,  2.,  2.,  2.]]])
>>>