Python: Stretch 2D array into 3D based on corresponding array with 3rd-dimension index

Viewed 212

Lets say I have some 2D array a = np.ones((3,3))

I want to stretch this array into 3 dimensions. I have array b, the same size as a, that provides the index in the 3rd dimension that each corresponding element in a needs to go too.

I also have 3D array c that is filled with NaNs. This is the array that the information from a should be put into. The remaining blank spaces that do not get "filled: can remain 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]]])

So, in the above example, I would want to end up with c[0,0,2] = 1.

I know I probably do this with some nested loops, but ideally I want this done in a more efficient/vectorized way.

3 Answers
Related