Python：基于具有第三维索引的相应数组将2D阵列拉伸成3D

Question

假设我有一些2D数组a = np.ones((3,3))

我想将这个数组拉伸成3维。 我有阵列b ，大小相同a ，其提供在所述第三尺寸，在每个对应的元素的索引a需要去了。

我也有充满NaN的3D阵列c 。 这是应该放入a的信息的数组。 剩余的空白空间没有被填满：可以保留NaNs。

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

我知道我可能会用一些嵌套循环来做这件事，但理想情况下我希望以更高效/矢量化的方式完成。

Answer 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_arrays来可视化 ）：

i, j, k = np.broadcast_arrays(np.arange(3)[:,None], np.arange(3), b)

print("first dimension index: ")
print(i)
print("second dimension index: ")
print(j)
print("third dimension index: ")
print(k)

first dimension index: 
[[0 0 0]
 [1 1 1]
 [2 2 2]]
second dimension index: 
[[0 1 2]
 [0 1 2]
 [0 1 2]]
third dimension index: 
[[2 2 2]
 [1 0 2]
 [1 0 0]]

现在，高级索引为（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

Answer 2

好的，所以这感觉就像一个完全黑客，但是诀窍是：

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

Answer 3

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

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一个维度。

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