0

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)
YGA
  • 9,546
  • 15
  • 47
  • 50
  • `bar[[[0],[1]], np.argsort(foo)]` seems to do what you want, though. – Alok Singhal Jan 07 '10 at 02:34
  • @Alok: yes, but the strategy seemed to fail on a more complicated example (I also have to say I find it somewhat cryptic). – YGA Jan 07 '10 at 21:15
  • @YGA: Can you tell us what the complicated example is? I will explain the method in more detail later today, if you would like. – Alok Singhal Jan 07 '10 at 21:21

3 Answers3

3

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.

Alok Singhal
  • 93,253
  • 21
  • 125
  • 158
2

A nice generic solution (with n rows to sort) is offered at this post, i.e.,

bar[np.arange(foo.shape[0])[:,None], np.argsort(foo)]
grovduck
  • 406
  • 5
  • 13
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))

kenm
  • 23,127
  • 2
  • 43
  • 62
  • (tried to add a comment, but it was misformatted because you can't have code in comments. See edit above) – YGA Jan 07 '10 at 01:16
  • Right you are - it takes based on the flattened dimensions, and 0, 1, and 2 all lie in the first row of that array. Here, this might work... – kenm Jan 07 '10 at 01:41
  • Does not work, always gives the 1st row of `bar` re-ordered differently. – Ηλίας Jun 15 '11 at 09:49