rearrange columns of multidimensional arrays efficiently

user_5197329 0
Viewed 58

I'm fairly certain this problem has an almost trivial solution, but for the life of me I can't seem to figure it out right now, nor find anything online, so bear with me please.

I have a 3D array of size (n x m x m), called v (think of it as n (m x m)-matrices.) And I wish to rearrange the columns in each matrix according to the indices I get from a sorting.

Hence I have the indices, idxs, (n x m), with which I wish to rearrange the matrix.

Now I can do the rearranging like this:

V = np.empty_like(v)
for i in range(v.shape[0]):
    V[i,:,np.arange(len(idxs[0]))] = v[i,:,idxs[i,:]]

But this is a very slow operation.

Does anyone have a better way of doing this rearranging?

linear-algebra numpy optimization python
1 Answers

Use the iterator as a ranged-array for advanced-indexing for a vectorized way -

I = np.arange(v.shape[0])[:,None]
V[I,:,np.arange(len(idxs[0]))] = v[I,:,idxs]

Another with simply indexing into v to directly get V -

V = v[np.arange(v.shape[0])[:,None],:,idxs].swapaxes(1,2)
user_32938810
Related