python - 使用 numpy.argsort 的输出对 numpy.ndarray 的每一列进行排序

Question

我想根据之前处理过的引用数组对 numpy 二维数组进行排序。我的想法是存储我的引用数组的 numpy.argsort 输出并使用它对其他数组进行排序:

In [13]: # my reference array
    ...: ref_arr = np.random.randint(10, 30, 12).reshape(3, 4)
Out[14]:
array([[12, 22, 12, 13],
       [28, 26, 21, 23],
       [24, 14, 16, 25]])

# desired output:
array([[12, 14, 12, 13],
       [24, 22, 16, 23],
       [28, 26, 21, 25]])

我尝试过的:

In [15]: # store the sorting matrix
    ...: sm = np.argsort(ref_arr, axis=0)
Out[16]:
array([[0, 2, 0, 0],
       [2, 0, 2, 1],
       [1, 1, 1, 2]])

但不幸的是，最后一步只适用于一维数组:

In [17]: ref_arr[sm]
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-17-48b785178465> in <module>()
----> 1 ref_arr[sm]

IndexError: index 3 is out of bounds for axis 0 with size 3

我找到了这个 Github issue这是针对此问题创建的，但不幸的是，通过提及我尝试的仅适用于一维数组的方法解决了该问题。 🙄

In a comment to this issue提到了一个与我的问题类似的例子。该代码段没有解决我的问题，因为它按行 而非列 对数组进行排序。但它暗示了我必须朝哪个方向移动...

a[np.arange(np.shape(a)[0])[:,np.newaxis], np.argsort(a)]

不幸的是，我对这个例子的理解不够，无法根据我的用例对其进行调整。也许有人可以在这里解释这种高级索引的工作原理？这可能使我能够自己解决问题，但我也不介意交 key 解决方案。 ;)

谢谢。

以防万一:我在 OS X 上使用 Python 3.6.1 和 numpy 1.12.1。

Answer 1

截至 2018 年 5 月，它可以使用 np.take_along_axis

完成

np.take_along_axis(ref_arr, sm, axis=0)
Out[25]: 
array([[10, 16, 15, 10],
       [13, 23, 24, 12],
       [28, 26, 28, 28]])