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
(Pdb)
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']],
dtype='|S1')
(Pdb)
source to share
bar.take(np.argsort(foo))
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.
Edit:
Try the following: bar.take(np.argsort(foo.ravel()).reshape(foo.shape))
source to share
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.
source to share
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)]
source to share