python矩阵返回最大值索引_python numpy：从矩阵中获取最小/最大n个值和索引的有效方式...

自从另一个答案的时候,NumPy已经添加了

numpy.partition和

numpy.argpartition的部分排序功能,可以在O(arr.size)时间或O(arr.size n * log(n))中执行此操作,如果需要元素按排序顺序排列.

numpy.partition(arr,n)返回数组的大小,其中第n个元素是数组排序时的大小.所有较小的元素都来自这个元素,所有更大的元素都来了.

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]