在3D数组中对2D数组的列进行单独排序

Question

我知道我们可以按照以下方式对2D numpy数组的列进行排序:

a = np.array([[1,4,7],
[3,1,5],
[9,5,8]])
a = a[:, a[1, :].argsort()]
Out: [[4,1,7],
[1,3,5],
[5,9,8]]

请注意，这确实是我想要的。现在对第二行(index=1)进行排序，并且对第0行和第2行中的值也进行相应的移动。也就是说，列的位置根据第1行的排序顺序而改变。

但现在我的问题:我没有一个2D数组，而是一个3D数组(即一个2D数组的数组)。

a = np.array([[[1,4,7],
[3,1,5],
[9,5,8]],
[[2,8,7],
[3,8,1],
[9,2,8]]])

我仍然希望对2D数组的列进行排序，分别基于它们各自行的值1。期望的结果是:

([[[4,1,7],
[1,3,5],
[5,9,8]],
[[7,2,8],
[1,3,8],
[8,9,2]]])

我尝试了以下操作，但结果不尽如人意:

a = a[:, :, a[: , 1, :].argsort()]

Answer 1

尝试np.take_along_axis:

np.take_along_axis(a,a[:,1].argsort()[:,None], axis=2)

:

array([[[4, 1, 7],
[1, 3, 5],
[5, 9, 8]],
[[7, 2, 8],
[1, 3, 8],
[8, 9, 2]]])

老实说，不要问我为什么它工作:-)

Answer 2

可以使用numpy.argsort和numpy.take_along_axis的组合:

idx = np.argsort(a, axis=2)
np.take_along_axis(a, idx[:,None,1], axis=2)

它的工作原理是从argsort获得排序顺序，然后只保留相关行(这里是1)，重塑以广播take_along_axis对所有其他行的操作。

输出:

array([[[4, 1, 7],
[1, 3, 5],
[5, 9, 8]],
[[7, 2, 8],
[1, 3, 8],
[8, 9, 2]]])