将列添加到更高维度

Question

让我们假设我有一个叫做A的数组

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

array([[[0., 0.],
    [0., 0.],
    [0., 0.]],
   [[0., 0.],
    [0., 0.],
    [0., 0.]],
   [[0., 0.],
    [0., 0.],
    [0., 0.]],
   [[0., 0.],
    [0., 0.],
    [0., 0.]]])

和另一个称为B的数组

B = np.arange(4)
array([0, 1, 2, 3])

而且我想在三维中进行类似串联的操作以得到此结果：

array([[[0., 0., 0.0],
    [0., 0., 0.0],
    [0., 0., 0.0]],
   [[0., 0., 1.0],
    [0., 0., 1.0],
    [0., 0., 1.0]],
   [[0., 0., 2.0],
    [0., 0., 2.0],
    [0., 0., 2.0]],
   [[0., 0., 3.0],
    [0., 0., 3.0],
    [0., 0., 3.0]]])

我尝试了一些实用的方法，但没有成功。我可以用什么方式做到这一点而不循环？

Answer 1

您可以使用np.ogrid定义一个开放式网格并使用它，以便分配在相应轴上广播的np.ogrid：

B

---

i,j,k=A.shape
i,j,k = np.ogrid[:i,:j,:k]
A[i,:,-1] = B[:,None,None,None]