NumPy提出了一种通过np.argmax获取数组最大值索引的方法。

我想要一个类似的东西,但返回N个最大值的索引。

例如,如果我有一个数组[1,3,2,4,5],那么nargmax(array, n=3)将返回对应于元素[5,4,3]的下标[4,3,1]。


当前回答

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

其他回答

我发现最直观的方法是使用np.unique。

其思想是,唯一方法返回输入值的索引。然后根据最大唯一值和索引,重新创建原始值的位置。

multi_max = [1,1,2,2,4,0,0,4]
uniques, idx = np.unique(multi_max, return_inverse=True)
print np.squeeze(np.argwhere(idx == np.argmax(uniques)))
>> [4 7]

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

较新的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)时间。

如果你碰巧在使用一个多维数组,那么你需要平展和解开索引:

def largest_indices(ary, n):
    """Returns the n largest indices from a numpy array."""
    flat = ary.flatten()
    indices = np.argpartition(flat, -n)[-n:]
    indices = indices[np.argsort(-flat[indices])]
    return np.unravel_index(indices, ary.shape)

例如:

>>> xs = np.sin(np.arange(9)).reshape((3, 3))
>>> xs
array([[ 0.        ,  0.84147098,  0.90929743],
       [ 0.14112001, -0.7568025 , -0.95892427],
       [-0.2794155 ,  0.6569866 ,  0.98935825]])
>>> largest_indices(xs, 3)
(array([2, 0, 0]), array([2, 2, 1]))
>>> xs[largest_indices(xs, 3)]
array([ 0.98935825,  0.90929743,  0.84147098])

我认为最省时的方法是手动遍历数组并保持k-size的min-heap,正如其他人所提到的那样。

我还想出了一个蛮力方法:

top_k_index_list = [ ]
for i in range(k):
    top_k_index_list.append(np.argmax(my_array))
    my_array[top_k_index_list[-1]] = -float('inf')

在使用argmax获取其索引后,将最大的元素设置为一个较大的负值。然后argmax的下一次调用将返回第二大的元素。 您可以记录这些元素的原始值,并在需要时恢复它们。