使用Python和NumPy从矩阵中获取最小/最大n值和索引的有效方法

Question

在给定NumPy矩阵(2D数组)的情况下,返回数组中的最小值/最大值n(及其索引)的有效方法是什么？

目前我有:

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

这比pyopencv生成我想要运行它的数组的图像模板匹配算法要长三倍,我认为这很愚蠢.

Answer 1

从另一个答案的时间开始,NumPy添加了部分排序numpy.partition和numpy.argpartition功能,允许您及时完成此O(arr.size)操作,或者O(arr.size+n*log(n))如果您需要按排序顺序排列元素.

numpy.partition(arr, n)返回一个数组的大小arr,其中n第i个元素是这将是什么,如果该阵列被分选.所有较小的元素都在该元素之前,所有更大的元素都在之后.

numpy.argpartition是numpy.partition为numpy.argsort是numpy.sort.

以下是如何使用这些函数来查找最小n元素的索引arr:

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

如果你需要按顺序索引,那么row_indices[0]最小元素的行也是如此,而不是只有一个n最小元素:

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]

Answer 2

由于NumPy中没有堆实现,可能你最好的猜测是对整个数组进行排序并采用最后的n元素:

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]

(与您的实现相比,这可能会以相反的顺序返回列表 - 我没有检查.)

在这个答案中给出了一个更有效的解决方案,适用于较新版本的NumPy .