python - 根据索引矩阵对 numpy 矩阵进行排序

Question

在 numpy 中，我们可以这样对数组进行排序:

>>> import numpy as np

>>> a = np.array([0, 100, 200])
>>> order = np.array([1, 2, 0])
>>> print(a[order])
[100 200   0]

但是，当“顺序”是矩阵时，这不起作用:

>>> A = np.array([    [0, 1, 2],
                      [3, 4, 5],
                      [6, 7, 8]])

>>> Ord = np.array([  [1, 0, 2],
                      [0, 2, 1],
                      [2, 1, 0]])

>>> print(A[Ord].shape)
(3, 3, 3)

我想像这样对“A”进行排序:

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

Answer 1

你可以使用 np.take_along_axis为此。

np.take_along_axis(A, Ord, axis=1)

输出

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

如文档中所述，它通常与生成索引的函数一起使用，例如 argsort。但我不确定这是否会推广到超过 2 个维度。