How to get indices of top-K values from a numpy array

Let suppose I have probabilities from a Pytorch or Keras predictions and result is with the softmax function

from scipy.special import softmax
probs = softmax(np.random.randn(20,10),1) # 20 instances and 10 class probabilities
probs

I want to find top-5 indices from this numpy array. All I want to do is to run a loop on the results something like:

for index in top_5_indices:
    if index in result:
        print('Found')

I'll get if my results are in top-5 results.

Pytorch has top-k function and I have seen numpy.argpartition but I have no idea how to get this done?

Solution

argpartition(a, k) function in numpy rearranges indices of input array a around the kth smallest element, so that all indices of smaller elements end up to the left, and all indices of bigger elements end up to the right. Not needing to sort all elements saves time: argpartition takes O(n) time, while argsort takes O(n log n) time.

So you can get the indices of 5 biggest elements like this:

np.argpartition(probs,-5)[-5:]