sorting via argsort - generalization to 2d matrices

Question

For sorting a numpy via argsort, we can do:

import numpy as np

x = np.random.rand(3)
x_sorted = x[np.argsort(x)]

I am looking for a numpy solution for the generalization to two or higher dimensions.

The indexing as in the 1d case won't work for 2d matrices.

Y = np.random.rand(4, 3)
sort_indices = np.argsort(Y)
#Y_sorted = Y[sort_indices] (what would that line be?)

Related: I am looking for a pure numpy answer that addresses the same problem as solved in this answer: https://stackoverflow.com/a/53700995/2272172

Answer 1

Use np.take_along_axis:

import numpy as np

np.random.seed(42)

x = np.random.rand(3)
x_sorted = x[np.argsort(x)]

Y = np.random.rand(4, 3)
sort_indices = np.argsort(Y)

print(np.take_along_axis(Y, sort_indices, axis=1))

print(np.array(list(map(lambda x, y: y[x], np.argsort(Y), Y))))  # the solution provided 

Output

[[0.15599452 0.15601864 0.59865848]
 [0.05808361 0.60111501 0.86617615]
 [0.02058449 0.70807258 0.96990985]
 [0.18182497 0.21233911 0.83244264]]

[[0.15599452 0.15601864 0.59865848]
 [0.05808361 0.60111501 0.86617615]
 [0.02058449 0.70807258 0.96990985]
 [0.18182497 0.21233911 0.83244264]]