Python：根据具有三维索引的对应数组将二维数组拉伸为三维数组

假设我有一些2D数组

a=np.one（（3,3））

我想把这个数组扩展成3维。我有一个数组

b

，大小与

a

相同，它在第三维提供了

a

中每个对应元素也需要的索引

我还有一个3D数组

c

，里面填充了NAN。这是

a

中的信息应放入的数组。未“填充”的剩余空格可以保留为NaN

>>> a = np.ones((3,3))
>>> b = np.random.randint(0,3,(3,3))
>>> c = np.empty((3,3,3))*np.nan
>>> 
>>> a
array([[ 1.,  1.,  1.],
       [ 1.,  1.,  1.],
       [ 1.,  1.,  1.]])
>>> b
array([[2, 2, 2],
       [1, 0, 2],
       [1, 0, 0]])
>>> c
array([[[ nan,  nan,  nan],
        [ nan,  nan,  nan],
        [ nan,  nan,  nan]],

       [[ nan,  nan,  nan],
        [ nan,  nan,  nan],
        [ nan,  nan,  nan]],

       [[ nan,  nan,  nan],
        [ nan,  nan,  nan],
        [ nan,  nan,  nan]]])

因此，在上面的例子中，我希望以c[0,0,2]=1结束

---

我知道我可能会对一些嵌套循环执行此操作，但理想情况下，我希望以更高效/矢量化的方式执行此操作。

您可以像这样使用奇特的索引，假设

b

中的最大值始终小于

c.shape[2]

：

n1, n2 = a.shape
c[np.arange(n1)[:,None], np.arange(n2), b] = a

c
#array([[[ nan,  nan,   1.],
#        [ nan,  nan,   1.],
#        [ nan,  nan,   1.]],

#       [[ nan,   1.,  nan],
#        [  1.,  nan,  nan],
#        [ nan,  nan,   1.]],

#       [[ nan,   1.,  nan],
#        [  1.,  nan,  nan],
#        [  1.,  nan,  nan]]])

---

这里我们对所有维度使用整数数组进行触发，三个数组按如下方式相互广播（这里我们使用

numpy.broacast_数组

将其可视化）：

现在，高级索引设置为（0，0，2），（0，1，2），（0，2，2）…，即从相同位置的每个数组中选取一个值，以形成元素的索引：

---

一些测试用例：

c[0,0,2]
#1.0

c[0,1,2]
#1.0

c[2,1,0]
#1.0

---

好的，这感觉像是一个完全的黑客，但有一个技巧：

a = np.ones((3,3))
b = np.array([[2, 2, 2],
              [1, 0, 2],
              [1, 0, 0]])
c = np.empty((3,3,3))*np.nan

z_coords = np.arange(3)

c[z_coords[None, None, :] == b[..., None]] = a.ravel()

我要做的是创建一个布尔索引数组，该数组对于我们要分配的索引是真的，然后分配这些索引

array([[[ nan,  nan,   1.],
        [ nan,  nan,   1.],
        [ nan,  nan,   1.]],

       [[ nan,   1.,  nan],
        [  1.,  nan,  nan],
        [ nan,  nan,   1.]],

       [[ nan,   1.,  nan],
        [  1.,  nan,  nan],
        [  1.,  nan,  nan]]])

---

一个较慢但可能更清晰的选择：

x, y = np.indices(c.shape[:2])
c[x, y, b] = a  # same as looping over c[x[i,j], y[i,j], b[i,j]] = a[i,j]

诀窍是生成3个具有相同形状的索引数组-每个

c

维度对应一个

---

公认的答案基本上是这样做的，但利用广播

这两种方法都有效，但这种方法比乔纳斯·阿德勒的更快。你能准确地解释一下这里发生了什么吗？@user3826115:B提供了第三维度的索引-

arange

s提供了其他两个维度的索引。

x, y = np.indices(c.shape[:2])
c[x, y, b] = a  # same as looping over c[x[i,j], y[i,j], b[i,j]] = a[i,j]