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

Posted 2023-02-23

【中文标题】使用 NumPy 从矩阵中获取最小/最大 n 值和索引的有效方法

【英文标题】：Efficient way to take the minimum/maximum n values and indices from a matrix using NumPy

【发布时间】：2011-08-14 00:05:48

【问题描述】：

在给定 NumPy 矩阵（二维数组）的情况下，返回数组中最小/最大 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 生成我要运行它的数组所用的图像模板匹配算法要长三倍，我认为这很愚蠢。

【问题讨论】：

n 与 len(arr) 的典型比率是多少？

@Paul: tiny.. 我正在查找模板与图像的匹配数，因此它是图像中像素数的匹配数，例如 20 到 150000

【参考方案1】：

由于 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 较新版本的更高效的解决方案。

【讨论】：

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

没有 NumPy 专家，但我们真的需要对在 O(n) 中简单完成的事情进行排序 (O(n log n)) 吗？我认为优点是排序是在 C 中完成的，而循环代码是由 python 解释器运行的？

@Voo：OP 算法的复杂度为O(m log n)，其中m 是数组中元素的数量，n 是要查找的最高元素的数量。我的答案中的算法是O(m log m)。 m 和 n 的这两个复杂性之间的因子在 OP 的上述评论中是 4，​​通过摆脱 Python 循环可以弥补这一点。正如 Paul 上面提到的，如果 n 真的很小，可能会有更好的选择。

@Voo：是的，复杂性并不是一切。在这种情况下，在 C 中完成此操作比我的要快很多（快 3 倍） - 并且足够让我不再担心它，但如果我需要更快的东西，我会回来更多。但是 - 你会如何在 O(n) 上简单地做到这一点？

NumPy 有 numpy.partition 和 numpy.argpartition，如果你需要的话，这会让你在 O(arr.size) 或 O(arr.size+n*log(n)) n 项按顺序排列。

【参考方案2】：

从另一个答案开始，NumPy 添加了numpy.partition 和numpy.argpartition 函数进行部分排序，允许您在O(arr.size) 时间执行此操作，如果您需要按排序顺序排列的元素，则可以在O(arr.size+n*log(n)) 执行此操作.

numpy.partition(arr, n) 返回一个大小为arr 的数组，其中nth 元素是对数组进行排序后的样子。所有较小的元素都在该元素之前，所有较大的元素都在之后。

numpy.argpartition 是 numpy.partition 就像 numpy.argsort 是 numpy.sort。

以下是如何使用这些函数来查找二维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]

【讨论】：

这个确切的代码给了我 ValueError: not enough values to unpack (expected 2, got 1)

@WayneFilkins：听起来你试图在一维数组而不是二维数组上使用它。 1D case 比较简单，但是不能将 2D case 代码扔到 1D 数组中。

【参考方案3】：

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

应用于任意轴。 K 返回排序结果和原始矩阵中的相应索引。

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