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

【问题标题】：Selecting indices for a 2d array in numpy在numpy中为二维数组选择索引
【发布时间】：2010-01-06 22:33:54
【问题描述】：

这在 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)

【问题讨论】：

bar[[[0],[1]], np.argsort(foo)] 似乎可以做你想做的事。
@Alok：是的，但是这个策略在一个更复杂的例子上似乎失败了（我也不得不说我觉得它有点神秘）。
@YGA：你能告诉我们这个复杂的例子是什么吗？如果您愿意，我将在今天晚些时候更详细地解释该方法。

标签： python numpy

【解决方案1】：

你想要

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

这是因为您需要两个索引来索引bar。 [[0], [1]] 是为了获得正确的广播。有关完全相同的问题和答案，请参阅this post on numpy-discussion 邮件列表。

【讨论】：

【解决方案2】：

at this post 提供了一个很好的通用解决方案（有 n 行要排序），即，

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

【讨论】：

【解决方案3】：

bar.take(np.argsort(foo)) 产生了你想要的输出，所以你应该看看它的documentation 以确保它确实按照你的想法做。

编辑：

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

【讨论】：

（尝试添加注释，但格式错误，因为您不能在 cmets 中拥有代码。请参阅上面的编辑）
是的，它基于展平的尺寸，0、1 和 2 都位于该数组的第一行。在这里，这可能有效...
不起作用，总是以不同的方式重新排序bar 的第一行。