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

Question

我希望能够根据之前处理过的参考数组对numpy 2D数组进行排序。 我的想法是将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问题，但不幸的是，它仅解决了我尝试的适用于一维数组的方法。

在该问题的评论中，提到了一个类似于我的问题的例子。这段代码片段没有解决我的问题，因为它是按行排序而不是按列排序。但是它给了我一个方向。

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

很遗憾，我并不理解这个例子，无法将其适应到我的使用场景上。或许有人能够解释一下这里的高级索引是如何工作的呢？这可能会使我自己解决问题，但当然，如果有现成的解决方案，我也不介意。

谢谢。

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

- wedi

---

不，对我来说这不起作用，因为它按行对数组进行排序。我想按列进行排序。但是我希望在理解该示例的工作原理后能够找到正确的解决方案。 - wedi

你的 argsort 调用没有按照你所说的轴进行排序。 - user2357112

@wedi：不，我的意思是你的argsort正在沿着行排序，而你说你想要另一个方向。 - user2357112

https://dev59.com/f1wY5IYBdhLWcg3wATpt#33141247 试图说明和解释 a[np.arange(np.shape...np.argsort(a)]。 - hpaulj

有一个numpy问题，关于使这种事情更容易。 - Eric

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]])

- vozman

Answer 2

基本上需要两个步骤：

1] 使用 axis=0 获取每列的 argsort 索引 -

sidx = ref_arr.argsort(axis=0)

2] 使用高级索引来使用sidx选择行，即索引到第一维，并使用另一个范围数组索引到第二维，以便覆盖所有列中的sidx索引 -

out = ref_arr[sidx, np.arange(sidx.shape[1])]

样例运行 -

In [185]: ref_arr
Out[185]: 
array([[12, 22, 12, 13],
       [28, 26, 21, 23],
       [24, 14, 16, 25]])

In [186]: sidx = ref_arr.argsort(axis=0)

In [187]: sidx
Out[187]: 
array([[0, 2, 0, 0],
       [2, 0, 2, 1],
       [1, 1, 1, 2]])

In [188]: ref_arr[sidx, np.arange(sidx.shape[1])]
Out[188]: 
array([[12, 14, 12, 13],
       [24, 22, 16, 23],
       [28, 26, 21, 25]])

- Divakar

---

谢谢！我已经得到了我的解决方案，并且通过查看它，我也理解了另一个例子。 :) - wedi