使用NumPy从矩阵获取最小/最大n值和索引的有效方法 [英] Efficient way to take the minimum/maximum n values and indices from a matrix using NumPy

问题描述

在给定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?

目前，我有:

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.

推荐答案

自另一个答案开始，NumPy已添加

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)返回大小为arr的数组，其中第n个元素是对数组进行排序后的元素.所有较小的元素都在该元素之前，而所有较大的元素都在之后.

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对应numpy.partition，numpy.argsort对应numpy.sort.

这是使用这些函数查找二维arr的最小n元素的索引的方式:

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)

如果您需要按顺序排列索引，那么row_indices[0]是最小元素的行，而不仅仅是n最小元素之一:

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]

这篇关于使用NumPy从矩阵获取最小/最大n值和索引的有效方法的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！