Reposition sub-elements of one array into different rows based on indices from another array?

Question

I am trying to adjust the order of an array that I created using the method in this post.

Basically depending on the row index of the array (where axis=0), the sub-element for that index has been removed (eg. if looking at index 1 then only elements 0 and 2 are present).

I need a way to position each element back into a row corresponding to its original index.

>>> a = array([[[222],[333]],
               [[111],[333]],
               [[111],[222]]])

Desired output:

>>> a = array([[[111],[111]],
               [[222],[222]],
               [[333],[333]]])

During my attempts I thought this might be getting close:

>>> n = a.shape[0]
>>> row_idx = np.nonzero(~np.eye(n,dtype=bool))[1].reshape(n,n-1) # dim0 indices

>>> print a.ravel() #flatten the array 
[222, 333, 111, 333, 111, 222]
>>> print row_idx #somehow use these indices to reposition corresponding elements
[1    2    0    2    0    1]

I am not sure how to use row_idx to align the elements into corresponding rows, or if that is even an efficient way to do this?

Thanks!

Answer 1

Use flattened row_idx to get the reversed-argsort and simply index into flattened a with it and finally reshape -

a.ravel()[row_idx.ravel().argsort()].reshape(a.shape)

Step-by-step run -

1) Flattened array and indices :

In [72]: a.ravel()
Out[72]: array([222, 333, 111, 333, 111, 222])

In [73]: row_idx.ravel()
Out[73]: array([1, 2, 0, 2, 0, 1])

2) We need to find the permuted version of flattened a, let's call it a0, such that a0[1] = 222, a0[2] = 333 and so on. To find those permuted indices, we need to use argsort on flattened row_idx. So, index into a with these for the final output. To get the output back in multi-dimensional format, we need that final reshape :

In [77]: a.ravel()[row_idx.ravel().argsort()]
Out[77]: array([111, 111, 222, 222, 333, 333])

In [78]: a.ravel()[row_idx.ravel().argsort()].reshape(a.shape)
Out[78]: 
array([[[111],
        [111]],

       [[222],
        [222]],

       [[333],
        [333]]])