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

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

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


当前回答

比较了编码的便捷性和速度

速度对我的需求很重要,所以我测试了这个问题的三个答案。

根据我的具体情况,对这三个答案中的代码进行了修改。

然后我比较了每种方法的速度。

编码智慧:

NPE的回答是最优雅的,也足够快地满足我的需求。 Fred foo的回答需要最多的重构来满足我的需求,但却是最快的。我选择了这个答案,因为尽管它需要更多的工作,但它并不太糟糕,并且具有显著的速度优势。 Off99555的回答是最优雅的,但也是最慢的。

测试和比较的完整代码

import numpy as np
import time
import random
import sys
from operator import itemgetter
from heapq import nlargest

''' Fake Data Setup '''
a1 = list(range(1000000))
random.shuffle(a1)
a1 = np.array(a1)

''' ################################################ '''
''' NPE's Answer Modified A Bit For My Case '''
t0 = time.time()
indices = np.flip(np.argsort(a1))[:5]
results = []
for index in indices:
    results.append((index, a1[index]))
t1 = time.time()
print("NPE's Answer:")
print(results)
print(t1 - t0)
print()

''' Fred Foos Answer Modified A Bit For My Case'''
t0 = time.time()
indices = np.argpartition(a1, -6)[-5:]
results = []
for index in indices:
    results.append((a1[index], index))
results.sort(reverse=True)
results = [(b, a) for a, b in results]
t1 = time.time()
print("Fred Foo's Answer:")
print(results)
print(t1 - t0)
print()

''' off99555's Answer - No Modification Needed For My Needs '''
t0 = time.time()
result = nlargest(5, enumerate(a1), itemgetter(1))
t1 = time.time()
print("off99555's Answer:")
print(result)
print(t1 - t0)

输出速度报告

肺水肿的回答是:

[(631934, 999999), (788104, 999998), (413003, 999997), (536514, 999996), (81029, 999995)]
0.1349949836730957

Fred Foo的回答:

[(631934, 999999), (788104, 999998), (413003, 999997), (536514, 999996), (81029, 999995)]
0.011161565780639648

off99555的回答是:

[(631934, 999999), (788104, 999998), (413003, 999997), (536514, 999996), (81029, 999995)]
0.439760684967041

其他回答

较新的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.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]

这将比完整排序更快,这取决于原始数组的大小和选择的大小:

>>> A = np.random.randint(0,10,10)
>>> A
array([5, 1, 5, 5, 2, 3, 2, 4, 1, 0])
>>> B = np.zeros(3, int)
>>> for i in xrange(3):
...     idx = np.argmax(A)
...     B[i]=idx; A[idx]=0 #something smaller than A.min()
...     
>>> B
array([0, 2, 3])

当然,这涉及到对原始数组的篡改。你可以修复(如果需要)通过复制或替换回原始值. ...对你的用例来说,哪个更便宜。

如果你正在处理nan和/或理解np有问题。试试pandas.DataFrame.sort_values。

import numpy as np
import pandas as pd    

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

df = pd.DataFrame(a, columns=['array'])
max_values = df['array'].sort_values(ascending=False, na_position='last')
ind = max_values[0:3].index.to_list()

这个例子给出了3个最大的非nan值的索引。可能效率很低,但易于阅读和定制。

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

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])