NumPy提出了一种通过np.argmax获取数组最大值索引的方法。
我想要一个类似的东西,但返回N个最大值的索引。
例如,如果我有一个数组[1,3,2,4,5],那么nargmax(array, n=3)将返回对应于元素[5,4,3]的下标[4,3,1]。
NumPy提出了一种通过np.argmax获取数组最大值索引的方法。
我想要一个类似的东西,但返回N个最大值的索引。
例如,如果我有一个数组[1,3,2,4,5],那么nargmax(array, n=3)将返回对应于元素[5,4,3]的下标[4,3,1]。
当前回答
我发现最直观的方法是使用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]
其他回答
我能想到的最简单的是:
>>> import numpy as np
>>> arr = np.array([1, 3, 2, 4, 5])
>>> arr.argsort()[-3:][::-1]
array([4, 3, 1])
这涉及到一个完整的数组。我想知道numpy是否提供了一种内置的方法来进行部分排序;到目前为止我还没有找到。
如果这个解决方案太慢(特别是对于小n),那么可能值得考虑用Cython编写一些东西。
这将比完整排序更快,这取决于原始数组的大小和选择的大小:
>>> 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])
当然,这涉及到对原始数组的篡改。你可以修复(如果需要)通过复制或替换回原始值. ...对你的用例来说,哪个更便宜。
简单的:
idx = (-arr).argsort()[:n]
其中n为最大值的个数。
使用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 ]])
Use:
>>> import heapq
>>> import numpy
>>> a = numpy.array([1, 3, 2, 4, 5])
>>> heapq.nlargest(3, range(len(a)), a.take)
[4, 3, 1]
对于常规的Python列表:
>>> a = [1, 3, 2, 4, 5]
>>> heapq.nlargest(3, range(len(a)), a.__getitem__)
[4, 3, 1]
如果使用Python 2,请使用xrange而不是range。
来源:堆队列算法