Claudiu Claudiu - 11 months ago 82
Python Question

python+numpy: efficient way to take the min/max n values and indices from a matrix

What's an efficient way, given a numpy matrix (2-d array), to return the min/max

values (along with their indices) in the array? Currently I have:

def n_max(arr, n):
res = [(0,(0,0))]*n
for y in xrange(len(arr)):
for x in xrange(len(arr[y])):
val = float(arr[y,x])
el = (val,(y,x))
i = bisect.bisect(res, el)
if i > 0:
res.insert(i, el)
del res[0]
return res

This takes 3x longer than the image template matching algorithm that
does to generate the array I want to run this on, and I figure that's silly.


Since there is no heap implementation in NumPy, probably your best guess is to sort the whole array and take the last n elements:

def n_max(arr, n):
    indices = arr.ravel().argsort()[-n:]
    indices = (numpy.unravel_index(i, arr.shape) for i in indices)
    return [(arr[i], i) for i in indices]

(This will probably return the list in reverse order compared to your implementation - did not check.)

Edit: A more efficient solution that works with newer versions of Numpy is given in this answer