根據各個行對 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]]])