通过argsort对二维矩阵的泛化进行排序

Question

要通过argsort对numpy进行排序，我们可以执行以下操作：

import numpy as np

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

我正在寻找一个numpy的解决方案推广到两个或更高的维度。你知道吗

一维情况下的索引不适用于二维矩阵。你知道吗

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

相关：我正在寻找一个纯粹的numpy答案，解决相同的问题，在这个答案：https://stackoverflow.com/a/53700995/2272172

Answer 1

使用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 

输出

[[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]]