在numpy中选择二维数组的索引

Question

在一维的情况下，这个方法效果很好：

# This will sort bar by the order of the values in foo
(Pdb) bar = np.array([1,2,3])
(Pdb) foo = np.array([5,4,6])
(Pdb) bar[np.argsort(foo)]
array([2, 1, 3])

但是在二维的情况下，我该怎么做呢？虽然argsort可以正常工作，但select就不行了：

(Pdb) foo = np.array([[5,4,6], [9,8,7]])
(Pdb) bar = np.array([[1,2,3], [1,2,3]])
(Pdb)  bar[np.argsort(foo)]
*** IndexError: index (2) out of range (0<=index<=1) in dimension 0
(Pdb) 

我本来期待这个输出是：

array([[2, 1, 3], [3, 2, 1]])

有没有人知道怎么做呢？

谢谢！

编辑：take()似乎可以做到，但它实际上只从第一行取元素（这真的让人困惑）。

如果我改变bar的值，你会看到：

(Pdb) bar = np.array([["1","2","3"], ["A", "B", "C"]])
(Pdb) bar.take(np.argsort(foo))
array([['2', '1', '3'],
       ['3', '2', '1']], 
      dtype='|S1')
(Pdb) 

Answer 1

在这个帖子中，提供了一个很不错的通用解决方案（适用于需要排序的n行数据），你可以查看一下这个链接：这里。

bar[np.arange(foo.shape[0])[:,None], np.argsort(foo)]

Answer 2

你想要的是

bar[[[0],[1]], np.argsort(foo)]

这是因为你需要两个索引来访问 bar。[[0], [1]] 是为了确保正确的广播操作。想了解更多，可以查看这个关于numpy讨论的帖子，里面有同样的问题和答案。

Answer 3

bar.take(np.argsort(foo)) 这个代码能产生你想要的结果，所以你应该去看看它的 文档，确认它确实能做到你想要的事情。

补充：

试试这个：bar.take(np.argsort(foo.ravel()).reshape(foo.shape))