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

Question

可以说我有一些2D阵列 a = np.ones((3,3))

我想将这个数组拉伸成3维.我有一个数组b,大小相同a,它提供了第三维中的索引,每个对应的元素也a需要去.

我也有c充满NaN的3D阵列.这是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