python - argsort for a multidimensional ndarray -

i'm trying indices sort multidimensional array last axis, e.g.

>>> = np.array([[3,1,2],[8,9,2]]) 

and i'd indices i such that,

>>> a[i] array([[1, 2, 3],        [2, 8, 9]]) 

based on documentation of numpy.argsort thought should this, i'm getting error:

>>> a[np.argsort(a)] indexerror: index 2 out of bounds axis 0 size 2 

---

edit: need rearrange other arrays of same shape (e.g. array b such a.shape == b.shape) in same way... that

>>> b = np.array([[0,5,4],[3,9,1]]) >>> b[i] array([[5,4,0],        [9,3,1]]) 

solution:

>>> a[np.arange(np.shape(a)[0])[:,np.newaxis], np.argsort(a)] array([[1, 2, 3],        [2, 8, 9]]) 

---

you got right, though wouldn't describe cheating indexing.

maybe make clearer:

in [544]: i=np.argsort(a,axis=1)  in [545]: out[545]:  array([[1, 2, 0],        [2, 0, 1]]) 

i order want, each row. is:

in [546]: a[0, i[0,:]] out[546]: array([1, 2, 3])  in [547]: a[1, i[1,:]] out[547]: array([2, 8, 9]) 

to both indexing steps @ once, have use 'column' index 1st dimension.

in [548]: a[[[0],[1]],i] out[548]:  array([[1, 2, 3],        [2, 8, 9]]) 

another array paired i is:

in [560]: j=np.array([[0,0,0],[1,1,1]])  in [561]: j out[561]:  array([[0, 0, 0],        [1, 1, 1]])  in [562]: a[j,i] out[562]:  array([[1, 2, 3],        [2, 8, 9]]) 

if i identifies column each element, j specifies row each element. [[0],[1]] column array works because can broadcasted against i.

i think of

np.array([[0],           [1]]) 

as 'short hand' j. define source row , column of each element of new array. work together, not sequentially.

the full mapping a new array is:

[a[0,1]  a[0,2]  a[0,0]  a[1,2]  a[1,0]  a[1,1]] 

---

def foo(a):     = np.argsort(a, axis=1)     return (np.arange(a.shape[0])[:,none], i)  in [61]: foo(a) out[61]:  (array([[0],         [1]]), array([[1, 2, 0],         [2, 0, 1]], dtype=int32)) in [62]: a[foo(a)] out[62]:  array([[1, 2, 3],        [2, 8, 9]])