使用NumPy从矩阵中获取最小/最大n个值及其索引的高效方法

Question

使用NumPy从矩阵中获取最小/最大n个值及其索引的高效方法

18

如何高效地从一个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生成我想要运行的数组所需的图像模板匹配算法需要三倍的时间，我认为这很傻。

- Claudiu

“n”与“len(arr)”之间的典型比率是多少？ - Paul

@Paul：小问题..我正在寻找模板匹配图像的数量，因此它是匹配数除以图像中的像素数，例如20到150000。 - Claudiu

3个回答

9

由于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，在此答案中给出。

- Sven Marnach

2

如果 n 很小，那么运行几次 argmax（每次删除最大值）可能会更快。 - Paul

1

我不是NumPy的专家，但我们真的需要对一些可以在O(n)中轻松完成的东西进行排序（O(n log n)）吗？我猜优势在于排序是由C完成的，而循环代码则由Python解释器运行？ - Voo

1

@Voo：OP算法的复杂度为O(m log n)，其中m是数组中元素的数量，n是要查找的最高元素的数量。我答案中的算法是O(m log m)。如OP上面的评论中所述，对于m和n之间的因子4，通过消除Python循环进行补偿更多。正如Paul上面指出的，如果n非常小，可能会有更好的选择。 - Sven Marnach

1

@Voo：是的，复杂度并不是一切。在这种情况下，用C语言完成比我的方法快得多（大约快3倍）- 而且足够快，以至于我不再需要担心它了，尽管如果我需要更快的速度，我会回来寻求更多帮助。但是，你如何轻松地在O(n)时间复杂度内完成呢？ - Claudiu

1

NumPy拥有numpy.partition和numpy.argpartition，可以让您以O(arr.size)或O(arr.size+n*log(n))的时间复杂度完成此操作，如果您需要按顺序获取n个项目。 - user2357112

显示剩余5条评论

0

我刚遇到同样的问题并解决了它。
以下是我的解决方案，通过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

- Arthur Chan

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- user2357112 · Accepted Answer

自另一个答案以来，NumPy已添加了numpy.partition和numpy.argpartition函数用于部分排序。如果您需要排序后的元素，则可以在O(arr.size)时间内完成，或者在O(arr.size+n*log(n))时间内完成。

numpy.partition(arr, n)返回一个大小与arr相同的数组，其中第n个元素是如果按顺序排列，则其可能的值。所有较小的元素都在该元素之前，而所有较大的元素都在之后。

numpy.argpartition类似于numpy.argsort和numpy.sort之间的关系，它是对numpy.partition的补充。

这是如何使用这些函数来查找二维数组arr中最小的n个元素的索引：

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]

一维情况简单得多：

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