这段代码适用于numpy 2D矩阵数组:

mat = np.array([[1, 3], [2, 5]]) # numpy matrix
 
n = 2  # n
n_largest_mat = np.sort(mat, axis=None)[-n:] # n_largest 
tf_n_largest = np.zeros((2,2), dtype=bool) # all false matrix
for x in n_largest_mat: 
  tf_n_largest = (tf_n_largest) | (mat == x) # true-false  

n_largest_elems = mat[tf_n_largest] # true-false indexing 

这将产生一个true-false的n_maximum矩阵索引，也可以从矩阵数组中提取n_maximum元素

如果你不关心第k大元素的顺序，你可以使用argpartition，它应该比通过argsort进行完整排序执行得更好。

K = 4 # We want the indices of the four largest values
a = np.array([0, 8, 0, 4, 5, 8, 8, 0, 4, 2])
np.argpartition(a,-K)[-K:]
array([4, 1, 5, 6])

这个问题值得肯定。

我运行了一些测试，看起来随着数组的大小和K值的增加，argpartition的性能优于argsort。

下面是查看最大元素及其位置的一个非常简单的方法。这里轴是定义域;对于2D情况，axis = 0表示列的最大数量，axis = 1表示行的最大数量。对于高维，这取决于你。

M = np.random.random((3, 4))
print(M)
print(M.max(axis=1), M.argmax(axis=1))

我能想到的最简单的是:

>>> import numpy as np
>>> arr = np.array([1, 3, 2, 4, 5])
>>> arr.argsort()[-3:][::-1]
array([4, 3, 1])

这涉及到一个完整的数组。我想知道numpy是否提供了一种内置的方法来进行部分排序;到目前为止我还没有找到。

如果这个解决方案太慢(特别是对于小n)，那么可能值得考虑用Cython编写一些东西。

Use:

def max_indices(arr, k):
    '''
    Returns the indices of the k first largest elements of arr
    (in descending order in values)
    '''
    assert k <= arr.size, 'k should be smaller or equal to the array size'
    arr_ = arr.astype(float)  # make a copy of arr
    max_idxs = []
    for _ in range(k):
        max_element = np.max(arr_)
        if np.isinf(max_element):
            break
        else:
            idx = np.where(arr_ == max_element)
        max_idxs.append(idx)
        arr_[idx] = -np.inf
    return max_idxs

它也适用于2D数组。例如,

In [0]: A = np.array([[ 0.51845014,  0.72528114],
                     [ 0.88421561,  0.18798661],
                     [ 0.89832036,  0.19448609],
                     [ 0.89832036,  0.19448609]])
In [1]: max_indices(A, 8)
Out[1]:
    [(array([2, 3], dtype=int64), array([0, 0], dtype=int64)),
     (array([1], dtype=int64), array([0], dtype=int64)),
     (array([0], dtype=int64), array([1], dtype=int64)),
     (array([0], dtype=int64), array([0], dtype=int64)),
     (array([2, 3], dtype=int64), array([1, 1], dtype=int64)),
     (array([1], dtype=int64), array([1], dtype=int64))]

In [2]: A[max_indices(A, 8)[0]][0]
Out[2]: array([ 0.89832036])

方法np。Argpartition只返回k个最大的索引，执行局部排序，比np快。当数组相当大时，Argsort(执行完全排序)。但是返回的索引不是升序或降序。让我们举个例子:

我们可以看到如果你想要一个严格的升序前k个指标，np。Argpartition不会返回你想要的。

除了在np后手动进行排序。argpartition，我的解决方案是使用PyTorch, torch。topk，一个神经网络构建工具，提供numpy类api，同时支持CPU和GPU。它和NumPy的MKL一样快，如果你需要大型矩阵/向量计算，它还提供了GPU的提升。

严格的上升/下降上k指数代码将是:

注意那个火炬。topk接受一个torch张量，并返回torch. tensor类型的topk值和topk索引。与np、torch类似。Topk还接受轴参数，以便处理多维数组/张量。