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

Question

这在1维上效果很好：

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

有什么线索怎么做？

谢谢！ / YGA

编辑： 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

你要

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

这是因为您需要两个索引来对bar进行索引。 [[0], [1]]用于获得正确的广播。 有关完全相同的问题和答案，请参见numpy-discussion邮件列表上的这篇文章 。

Answer 2

这篇文章提供了一个很好的通用解决方案（可以排序n行），即

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

Answer 3

bar.take(np.argsort(foo))产生了所需的输出，因此您应该查看其文档以确保其实际执行了您想要的操作。

编辑：

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