根据各个行对 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]]

请注意，这确实是我想要的。 第二行（索引 = 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 行的值对 2D 数组的列进行单独排序。所需的结果是：

([[[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在所有其他行上广播take_along_axis的操作。

输出：

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

       [[7, 2, 8],
        [1, 3, 8],
        [8, 9, 2]]])