python矩阵返回最大值索引_使用NumPy从矩阵获取最小/最大n值和索引的有效方法...

What's an efficient way, given a NumPy matrix (2D array), to return the minimum/maximum n 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 three times longer than the image template matching algorithm that pyopencv does to generate the array I want to run this on, and I figure that's silly.

解决方案

Since the time of the other answer, NumPy has added the numpy.partition and numpy.argpartition functions for partial sorting, allowing you to do this in O(arr.size) time, or O(arr.size+n*log(n)) if you need the elements in sorted order.

numpy.partition(arr, n) returns an array the size of arr where the nth element is what it would be if the array were sorted. All smaller elements come before that element and all greater elements come afterward.

numpy.argpartition is to numpy.partition as numpy.argsort is to numpy.sort.

Here's how you would use these functions to find the indices of the minimum n elements of a two-dimensional arr:

flat_indices = numpy.argpartition(arr.ravel(), n-1)[:n]

row_indices, col_indices = numpy.unravel_index(flat_indices, arr.shape)

And if you need the indices in order, so row_indices[0] is the row of the minimum element instead of just one of the n minimum elements:

min_elements = arr[row_indices, col_indices]

min_elements_order = numpy.argsort(min_elements)

row_indices, col_indices = row_indices[min_elements_order], col_indices[min_elements_order]

The 1D case is a lot simpler:

# Unordered:

indices = numpy.argpartition(arr, n-1)[:n]

# Extra code if you need the indices in order:

min_elements = arr[indices]

min_elements_order = numpy.argsort(min_elements)

ordered_indices = indices[min_elements_order]