Sorting numpy array according to index provided by another array

29 Views Asked by At

A = [5,7,6] and I = [2,0,1].

I want to reorganize A such that I get A_new = [7,6,5]

However, when I use A_new = numpy.take_along_axis(A,I,1), I get A_new = [6,5,7].

In my case, index of 5 should be 2 instead of index of 6 being 2. Please let me know how this can be done without using for loops.

However, when I use A_new = numpy.take_along_axis(A,I,1), I get A_new = [6,5,7].

1

There are 1 best solutions below

0
hpaulj On BEST ANSWER

The 'trick' is to use I to assign values to a 'black' array:

In [192]: A = [5,7,6];I = [2,0,1]

In [193]: res=np.zeros_like(A); res[I] = A

In [194]: res
Out[194]: array([7, 6, 5])

Using I directly as index of A doesn't do what you want (but often does for a different definition of I). Comments have suggested argsort:

In [195]: np.array(A)[I]
Out[195]: array([6, 5, 7])

In [196]: np.argsort(I)
Out[196]: array([1, 2, 0], dtype=int64)

In [197]: np.array(A)[np.argsort(I)]
Out[197]: array([7, 6, 5])

argsort in effect converts I to one that does index A correctly.

In your definition of I, you want A[0] to be in the res[2] slot of the result. It's the difference between saying I want to move the 0'th item to the 2'th slot, versus I want 0's result to be 1'st' item.