Selecting indices for a 2d array in numpy

Question

Hi Folks,

This works quite well in 1 dimension:

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

But how do I do that in two dimensions? Argsort works nicely, but the select no longer works:

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

I would expect this to output:

array([[2, 1, 3], [3, 2, 1]])

Any clue how to do it?

Thanks! /YGA

Edit: take() would seem to do the right thing, but it really only takes elements from the first row (super confusing).

You can see that if I change the values of 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.take(np.argsort(foo)) produced your desired output, so you should take a look at its documentation to make sure it actually does what you think you want.

Edit:

Try this: bar.take(np.argsort(foo.ravel()).reshape(foo.shape))

Answer 2

You want

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

This is because you need two indices to index bar. The [[0], [1]] is to get correct broadcasting. See this post on numpy-discussion mailing list for exactly the same question and the answer.