Malintha Malintha - 1 month ago 10
Python Question

Get top N elements from array in descending order

I have an array

[2, 4, 2, 1, 3, 2, 2]
and I want to get top 3 value indexes in descending order. So my output should be [1,4,0] which are the indexes of top 3 elements of the array in descending order. (Similar item indexes are tolerated). I am using bottelneck which is doing my job partially.

import bottleneck
import numpy as np
a = np.array([2, 4, 2, 1, 3, 2, 2])
b = np.array(bottleneck.argpartsort(-a, 3)[:3])
print(b)


which gives
[4 1 5]
which are the correct top 3(n) indexes of the array but not in the descending order. Is there any built-in function to adjust this output to provide descending order

Answer

You can use np.argsort

import numpy as np
a = np.array([2, 4, 2, 1, 3, 2, 2])
b = np.argsort(a)[-3:]
b = b[::-1]

Note: The last two lines can further be simplified to (although with not much computational advantage)

b = np.argsort(a)[-3:][::-1]
# or, as juanpa suggested
b = np.argsort(a)[-1:-4:-1]

All give the following value of b

array([1, 4, 6])