Choosing indices for a 2d array in numpy

This works 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 this in two dimensions? Argsort works nicely, but selection doesn't work anymore:

(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


I expect this to output:

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


Any clue how to do this?

Thank! / Yga

Edit: take()

seems to do things right, but it actually only accepts elements from the first line (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']], 



source to share

3 answers


produced the desired output, so you should look at its documentation to make sure it actually does what you think you want it to do.


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



Do you want to

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


This is because you need two indexes to index bar

. [[0], [1]]

is correct broadcasting. See this post on the numpy-discussion posts list for exactly the same question and answer.



A good one-stop solution (with n lines to sort) is suggested in this post , i.e.

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




All Articles