Create list of index of array sorted from largest to smallest

Question

I have an array y_probability which is 2 dimensional. It contains the probability that an instance belongs to each class. I would like to create a list of the indexes of the y_probabiltiy array ordered from largest to smallest on the first column. How can I do this?

Sample data:

y_probability = np.array([[0.3,0.7],[0.5,0.5] ,[0.2,0.8], [0.1,0.9]])

Desired output:

index = [2,1,3,4]

Answer 1

You can use argsort on the first column (note the index in python is 0 based, so the result will be one less than what you expect), then reverse the result with [::-1] since argsort returns index that sorts array in ascending order:

y_probability[:,0].argsort()[::-1]
# array([1, 0, 2, 3])