python - 根据各自的行分别对 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]]])

我仍然想根据各自行 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]]])

老实说，不要问我为什么它有效:-)