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

(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[[[0],[1]], np.argsort(foo)]

Это потому, что вам нужно два индекса для индексации bar. [[0], [1]] - для правильного вещания. См. этот пост на numpy-обсуждение список рассылки для того же вопроса и ответа.

Answer 2

Хорошее общее решение (с n строками для сортировки) предлагается на этом посту , т.е.

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

Answer 3

bar.take(np.argsort(foo)) произвел желаемый результат, поэтому вам следует взглянуть на его документацию , чтобы убедиться, что он действительно выполняет то, что вы хотите.

Edit:

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