如何在NumPy数组中获得N个最大值的索引？

Question

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

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

例如,如果我有一个数组,[1, 3, 2, 4, 5],function(array, n=3)将返回的索引[4, 3, 1]相对应的元素[5, 4, 3].

Answer 1

较新的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])
>>> a[ind]
array([4, 9, 6, 9])

与argsort此不同,此函数在最坏的情况下以线性时间运行,但返回的索引未排序,从评估结果可以看出a[ind].如果您也需要,请在之后对其进行排序:

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

以这种方式按排序顺序获取top- k元素需要O(n + k log k)时间.

Answer 2

我能想到的最简单的是:

In [1]: import numpy as np

In [2]: arr = np.array([1, 3, 2, 4, 5])

In [3]: arr.argsort()[-3:][::-1]
Out[3]: array([4, 3, 1])

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

如果这个解决方案太慢(特别是对于小型n),那么在Cython中编写代码可能是值得的.

Answer 3

更简单:

idx = (-arr).argsort()[:n]

其中n是最大值的数量.

Answer 4

使用:

>>> 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.

来源:heapq - 堆队列算法

Answer 5

如果您正在使用多维数组,那么您将需要展平并解开索引:

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

Answer 6

如果你不关心你可以使用的第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])

积分转到这个问题.

我跑了几个测试,它看起来像argpartition性能优于argsort作为数组的大小和K寿命值.

Answer 7

三个答案比较编码的简易性和速度

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

根据我的具体情况，根据需要修改了这三个答案中的代码。

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

编码明智：

1. NPE 的答案是下一个最优雅且足够快的答案，可以满足我的需求。
2. Fred Foos 的回答需要对我的需求进行最多的重构，但速度最快。我选择了这个答案，因为即使需要做更多的工作，它也还不错，并且具有显着的速度优势。
3. 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)

输出速度报告

NPE's Answer:
[(631934, 999999), (788104, 999998), (413003, 999997), (536514, 999996), (81029, 999995)]
0.1349949836730957

Fred Foo's Answer:
[(631934, 999999), (788104, 999998), (413003, 999997), (536514, 999996), (81029, 999995)]
0.011161565780639648

off99555's Answer:
[(631934, 999999), (788104, 999998), (413003, 999997), (536514, 999996), (81029, 999995)]
0.439760684967041

Answer 8

对于多维数组,您可以使用axis关键字以沿预期轴应用分区.

# For a 2D array
indices = np.argpartition(arr, -N, axis=1)[:, -N:]

并抓住物品:

x = arr.shape[0]
arr[np.repeat(np.arange(x), N), indices.ravel()].reshape(x, N)

但请注意,这不会返回排序结果.在这种情况下,您可以np.argsort()沿预期的轴使用:

indices = np.argsort(arr, axis=1)[:, -N:]

# Result
x = arr.shape[0]
arr[np.repeat(np.arange(x), N), indices.ravel()].reshape(x, N)

这是一个例子:

In [42]: a = np.random.randint(0, 20, (10, 10))

In [44]: a
Out[44]:
array([[ 7, 11, 12,  0,  2,  3,  4, 10,  6, 10],
       [16, 16,  4,  3, 18,  5, 10,  4, 14,  9],
       [ 2,  9, 15, 12, 18,  3, 13, 11,  5, 10],
       [14,  0,  9, 11,  1,  4,  9, 19, 18, 12],
       [ 0, 10,  5, 15,  9, 18,  5,  2, 16, 19],
       [14, 19,  3, 11, 13, 11, 13, 11,  1, 14],
       [ 7, 15, 18,  6,  5, 13,  1,  7,  9, 19],
       [11, 17, 11, 16, 14,  3, 16,  1, 12, 19],
       [ 2,  4, 14,  8,  6,  9, 14,  9,  1,  5],
       [ 1, 10, 15,  0,  1,  9, 18,  2,  2, 12]])

In [45]: np.argpartition(a, np.argmin(a, axis=0))[:, 1:] # 1 is because the first item is the minimum one.
Out[45]:
array([[4, 5, 6, 8, 0, 7, 9, 1, 2],
       [2, 7, 5, 9, 6, 8, 1, 0, 4],
       [5, 8, 1, 9, 7, 3, 6, 2, 4],
       [4, 5, 2, 6, 3, 9, 0, 8, 7],
       [7, 2, 6, 4, 1, 3, 8, 5, 9],
       [2, 3, 5, 7, 6, 4, 0, 9, 1],
       [4, 3, 0, 7, 8, 5, 1, 2, 9],
       [5, 2, 0, 8, 4, 6, 3, 1, 9],
       [0, 1, 9, 4, 3, 7, 5, 2, 6],
       [0, 4, 7, 8, 5, 1, 9, 2, 6]])

In [46]: np.argpartition(a, np.argmin(a, axis=0))[:, -3:]
Out[46]:
array([[9, 1, 2],
       [1, 0, 4],
       [6, 2, 4],
       [0, 8, 7],
       [8, 5, 9],
       [0, 9, 1],
       [1, 2, 9],
       [3, 1, 9],
       [5, 2, 6],
       [9, 2, 6]])

In [89]: a[np.repeat(np.arange(x), 3), ind.ravel()].reshape(x, 3)
Out[89]:
array([[10, 11, 12],
       [16, 16, 18],
       [13, 15, 18],
       [14, 18, 19],
       [16, 18, 19],
       [14, 14, 19],
       [15, 18, 19],
       [16, 17, 19],
       [ 9, 14, 14],
       [12, 15, 18]])

Answer 9

这将比完整排序更快，具体取决于原始数组的大小和选择的大小：

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

当然，它涉及篡改您的原始数组。您可以通过复制或替换原始值来修复（如果需要）。...以您的用例而言更便宜的为准。

Answer 10

方法np.argpartition只返回 k 个最大的索引，执行局部排序，并且比np.argsort数组非常大时（执行完整排序）更快。但是返回的索引不是按升序/降序排列的。让我们用一个例子说：

我们可以看到，如果你想要一个严格的升序前 k 个索引，np.argpartition不会返回你想要的。

除了在 np.argpartition 之后手动进行排序，我的解决方案是使用 PyTorch，torch.topk一个神经网络构建工具，提供类似 NumPy 的 API，同时支持 CPU 和 GPU。它与带有 MKL 的 NumPy 一样快，如果您需要大型矩阵/向量计算，它可以提供 GPU 提升。

严格的上升/下降前 k 个索引代码将是：

请注意，torch.topk接受一个火炬张量，并在 type 中返回前 k 个值和前 k 个索引torch.Tensor。与 np 类似，torch.topk 也接受轴参数，以便您可以处理多维数组/张量。