Have numpy argsort return an array of 2d indices?
Use np.take_along_axis
to slices with your indices. Assuming you need to sort row-wise:
>>> arr = np.array([[5, 2, 4],
[3, 3, 3],
[6, 1, 2]])
>>> ind = arr.argsort(axis=-1)
>>> ind
array([[1, 2, 0],
[0, 1, 2],
[1, 2, 0]])
>>> np.take_along_axis(arr, ind, axis=-1)
array([[2, 4, 5],
[3, 3, 3],
[1, 2, 6]])
Apply numpy.argsort
on flattened array and then unravel the indices back to (3, 3) shape:
>>> arr = np.array([[5, 2, 4],
[3, 3, 3],
[6, 1, 2]])
>>> np.dstack(np.unravel_index(np.argsort(arr.ravel()), (3, 3)))
array([[[2, 1],
[0, 1],
[2, 2],
[1, 0],
[1, 1],
[1, 2],
[0, 2],
[0, 0],
[2, 0]]])
From the documentation on numpy.argsort:
ind = np.unravel_index(np.argsort(x, axis=None), x.shape)
Indices of the sorted elements of a N-dimensional array.
An example:
>>> x = np.array([[0, 3], [2, 2]])
>>> x
array([[0, 3],
[2, 2]])
>>> ind = np.unravel_index(np.argsort(x, axis=None), x.shape)
>>> ind # a tuple of arrays containing the indexes
(array([0, 1, 1, 0]), array([0, 0, 1, 1]))
>>> x[ind] # same as np.sort(x, axis=None)
array([0, 2, 2, 3])