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元素是数组排序时的元素。所有较小的元素都在该元素之前，所有更大的元素都在之后。

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]

1D的情况要简单得多：

# 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]

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]

（与你的实现相比，这可能会以相反的顺序返回列表 - 我没有检查。）

this answer中提供了一个更有效的解决方案，适用于较新版本的NumPy。

Answer 3

我刚刚遇到了完全相同的问题并解决了。
这是我的解决方案，包装了np.argpartition：

应用于任意轴。
当K << array.shape [axis]，o（N）时保持高速。
返回排序后的结果和原始矩阵中的相应索引。

def get_sorted_smallest_K(array, K, axis=-1):
    # Find the least K values of array along the given axis. 
    # Only efficient when K << array.shape[axis].
    # Return:
    #   top_sorted_scores: np.array. The least K values.
    #   top_sorted_indexs: np.array. The least K indexs of original input array.
    
    partition_index = np.take(np.argpartition(array, K, axis), range(0, K), axis)
    top_scores = np.take_along_axis(array, partition_index, axis)
    sorted_index = np.argsort(top_scores, axis=axis)
    top_sorted_scores = np.take_along_axis(top_scores, sorted_index, axis)
    top_sorted_indexs = np.take_along_axis(partition_index, sorted_index, axis)
    return top_sorted_scores, top_sorted_indexs

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

3 个答案: