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