Search code examples
pythonnumpybroadcasting

How to efficiently order a numpy matrix


I have this numpy array

matrix = np.array([[ 0.8,  0.2,  0.1],
   [ 1. ,  0. ,  0. ],
   [ 0. ,  0. ,  1. ]])

and I would like to filter to return, for each row of matrix the indices in decreasing value order.

For example, this would be

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

I know I could use np.argsort, but this doesn't seem to be returning the right output. I tried changing the axis to different values, but that doesn't help either.


Solution

  • Probably the easiest way to get your desired output would be:

    (-matrix).argsort(axis=1)
    # array([[0, 1, 2],
    #        [0, 1, 2],
    #        [2, 0, 1]])