Выбор индексов для 2d массива в 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 работает хорошо, но выбор больше не работает:

(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() может показаться правильным, но на самом деле это занимает только элементы из первой строки (супер запутанной).

Вы можете видеть, что если я изменить значение бара:

(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.take(np.argsort(foo)) произвел свой желаемый результат, так что вы должны смотреть на его documentation, чтобы убедиться, что он на самом деле делает то, что вы думаете, что вы хотите.

Edit:

Попробуйте это: bar.take(np.argsort(foo.ravel()).reshape(foo.shape))

Answer 2

Вы хотите

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

Это потому, что вам нужно два индекса к индексу bar. [[0], [1]] - это правильное вещание. См. this post on numpy-discussion список рассылки за тот же вопрос и ответ.

Answer 3

Хороший общее решение (с п строк для сортировки) предлагается at this post, т.е.

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