The argsort() function returns a matrix of indices that can be used to index the original array so that the result would match the sort() result.

Is there a way to apply those indices? I have two arrays, one is the array used for obtaining the sort order, and another is some associated data.

I would like to compute assoc_data[array1.argsort()] but that doesn't seem to work.

Here's an example:

z=array([1,2,3,4,5,6,7])
z2=array([z,z*z-7])
i=z2.argsort()

z2=array([[ 1,  2,  3,  4,  5,  6,  7],[-6, -3,  2,  9, 18, 29, 42]])
i =array([[1, 1, 1, 0, 0, 0, 0],[0, 0, 0, 1, 1, 1, 1]])

I would like to apply i to z2 (or another array with associated data) but I'm not sure how to do so.

This is probably overkill, but this will work in the nd case:

import numpy as np
axis = 0
index = list(np.ix_(*[np.arange(i) for i in z2.shape]))
index[axis] = z2.argsort(axis)
z2[index]# Or if you only need the 3d case you can use np.ogrid.axis = 0
index = np.ogrid[:z2.shape[0], :z2.shape[1], :z2.shape[2]]
index[axis] = z2.argsort(axis)
z2[index]