I have need the N minimum (index) values in a numpy array I have need the N minimum (index) values in a numpy array python python

I have need the N minimum (index) values in a numpy array


If you call

arr.argsort()[:3]

It will give you the indices of the 3 smallest elements.

array([0, 2, 1], dtype=int64)

So, for n, you should call

arr.argsort()[:n]


Since this question was posted, numpy has updated to include a faster way of selecting the smallest elements from an array using argpartition. It was first included in Numpy 1.8.

Using snarly's answer as inspiration, we can quickly find the k=3 smallest elements:

In [1]: import numpy as npIn [2]: arr = np.array([1, 3, 2, 4, 5])In [3]: k = 3In [4]: ind = np.argpartition(arr, k)[:k]In [5]: indOut[5]: array([0, 2, 1])In [6]: arr[ind]Out[6]: array([1, 2, 3])

This will run in O(n) time because it does not need to do a full sort. If you need your answers sorted (Note: in this case the output array was in sorted order but that is not guaranteed) you can sort the output:

In [7]: sorted(arr[ind])Out[7]: array([1, 2, 3])

This runs on O(n + k log k) because the sorting takes place on the smalleroutput list.


I don't guarantee that this will be faster, but a better algorithm would rely on heapq.

import heapqindices = heapq.nsmallest(10,np.nditer(arr),key=arr.__getitem__)

This should work in approximately O(N) operations whereas using argsort would take O(NlogN) operations. However, the other is pushed into highly optimized C, so it might still perform better. To know for sure, you'd need to run some tests on your actual data.