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

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

使用argpartition的向量化2D实现:

k = 3
probas = np.array([
    [.6, .1, .15, .15],
    [.1, .6, .15, .15],
    [.3, .1, .6, 0],
])

k_indices = np.argpartition(-probas, k-1, axis=-1)[:, :k]

# adjust indices to apply in flat array
adjuster = np.arange(probas.shape[0]) * probas.shape[1]
adjuster = np.broadcast_to(adjuster[:, None], k_indices.shape)
k_indices_flat = k_indices + adjuster

k_values = probas.flatten()[k_indices_flat]

# k_indices:
# array([[0, 2, 3],
#        [1, 2, 3],
#        [2, 0, 1]])
# k_values:
# array([[0.6 , 0.15, 0.15],
#        [0.6 , 0.15, 0.15],
#       [0.6 , 0.3 , 0.1 ]])

当top_k<<axis_length时，它优于argsort。

import numpy as np

def get_sorted_top_k(array, top_k=1, axis=-1, reverse=False):
    if reverse:
        axis_length = array.shape[axis]
        partition_index = np.take(np.argpartition(array, kth=-top_k, axis=axis),
                                  range(axis_length - top_k, axis_length), axis)
    else:
        partition_index = np.take(np.argpartition(array, kth=top_k, axis=axis), range(0, top_k), axis)
    top_scores = np.take_along_axis(array, partition_index, axis)
    # resort partition
    sorted_index = np.argsort(top_scores, axis=axis)
    if reverse:
        sorted_index = np.flip(sorted_index, axis=axis)
    top_sorted_scores = np.take_along_axis(top_scores, sorted_index, axis)
    top_sorted_indexes = np.take_along_axis(partition_index, sorted_index, axis)
    return top_sorted_scores, top_sorted_indexes

if __name__ == "__main__":
    import time
    from sklearn.metrics.pairwise import cosine_similarity

    x = np.random.rand(10, 128)
    y = np.random.rand(1000000, 128)
    z = cosine_similarity(x, y)
    start_time = time.time()
    sorted_index_1 = get_sorted_top_k(z, top_k=3, axis=1, reverse=True)[1]
    print(time.time() - start_time)

我能想到的最简单的是:

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

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

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

较新的NumPy版本(1.8及以上)有一个名为argpartition的函数。要得到四个最大元素的索引，请执行以下操作

>>> a = np.array([9, 4, 4, 3, 3, 9, 0, 4, 6, 0])
>>> a
array([9, 4, 4, 3, 3, 9, 0, 4, 6, 0])

>>> ind = np.argpartition(a, -4)[-4:]
>>> ind
array([1, 5, 8, 0])

>>> top4 = a[ind]
>>> top4
array([4, 9, 6, 9])

与argsort不同，这个函数在最坏的情况下以线性时间运行，但返回的索引没有排序，这可以从求值a[ind]的结果中看出。如果你也需要，那就把它们分类:

>>> ind[np.argsort(a[ind])]
array([1, 8, 5, 0])

以这种方式获得排在前k位的元素需要O(n + k log k)时间。

方法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还接受轴参数，以便处理多维数组/张量。