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]。
当前回答
如果你不关心第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。
其他回答
Use:
from operator import itemgetter
from heapq import nlargest
result = nlargest(N, enumerate(your_list), itemgetter(1))
现在,结果列表将包含N个元组(index, value),其中value是最大的。
下面是查看最大元素及其位置的一个非常简单的方法。这里轴是定义域;对于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 ]])
如果你不关心第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。
简单的:
idx = (-arr).argsort()[:n]
其中n为最大值的个数。