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

NumPy提出了一种通过

np.argmax

获取数组最大值索引的方法

我想要一个类似的东西，但是返回

N

最大值的索引

---

例如，如果我有一个数组，

[1,3,2,4,5]

，

函数（数组，n=3）

将返回与元素

[5,4,3]

相对应的索引

[4,3,1]

我能想到的最简单的方法是：

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

），则可能值得考虑在中对某些内容进行编码。

具有部分排序功能，如果仅为了获得n个最大值而对整个数组进行排序的开销太大

---

我对这个模块一无所知；我刚在谷歌上搜索了

numpy部分排序

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

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

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

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

使用：

>>> import heapq
>>> import numpy
>>> a = numpy.array([1, 3, 2, 4, 5])
>>> heapq.nlargest(3, range(len(a)), a.take)
[4, 3, 1]

from operator import itemgetter
from heapq import nlargest
result = nlargest(N, enumerate(your_list), itemgetter(1))

def max_indices(arr, k):
    '''
    Returns the indices of the k first largest elements of arr
    (in descending order in values)
    '''
    assert k <= arr.size, 'k should be smaller or equal to the array size'
    arr_ = arr.astype(float)  # make a copy of arr
    max_idxs = []
    for _ in range(k):
        max_element = np.max(arr_)
        if np.isinf(max_element):
            break
        else:
            idx = np.where(arr_ == max_element)
        max_idxs.append(idx)
        arr_[idx] = -np.inf
    return max_idxs

对于常规Python列表：

>>> a = [1, 3, 2, 4, 5]
>>> heapq.nlargest(3, range(len(a)), a.__getitem__)
[4, 3, 1]

如果使用Python2，请使用

xrange

而不是

range

来源：

较新的NumPy版本（1.8及以上）有一个为此调用的函数。要获得四个最大元素的索引，请执行以下操作

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

以这种方式获得按排序顺序排列的前k个元素需要O（n+k log k）时间。

更简单的是：

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

其中n是最大值的数目。

使用：

>>> import heapq
>>> import numpy
>>> a = numpy.array([1, 3, 2, 4, 5])
>>> heapq.nlargest(3, range(len(a)), a.take)
[4, 3, 1]

from operator import itemgetter
from heapq import nlargest
result = nlargest(N, enumerate(your_list), itemgetter(1))

def max_indices(arr, k):
    '''
    Returns the indices of the k first largest elements of arr
    (in descending order in values)
    '''
    assert k <= arr.size, 'k should be smaller or equal to the array size'
    arr_ = arr.astype(float)  # make a copy of arr
    max_idxs = []
    for _ in range(k):
        max_element = np.max(arr_)
        if np.isinf(max_element):
            break
        else:
            idx = np.where(arr_ == max_element)
        max_idxs.append(idx)
        arr_[idx] = -np.inf
    return max_idxs

---

现在，

结果

列表将包含N元组（

索引

，

值

），其中

值

最大化。

如果您不关心可以使用的第K个最大元素的顺序，它的性能应该比通过

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

。

如果您使用的是多维数组，则需要展平和分解索引：

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

---

对于多维数组，可以使用

轴

关键字沿预期轴应用分区

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

---

我发现使用

np.unique

最直观

其思想是，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

（执行完全排序）更快。但是返回的索引不是按升序/降序排列的。让我们举一个例子：

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

我们可以看到，如果您想要一个严格的升序top k索引，

np.argpartition

将不会返回您想要的

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

严格的上升/下降顶部k索引代码为：

请注意，它接受torch张量，并在类型

torch.tensor

中返回top k值和top k索引。与np类似，torch.topk也接受轴参数，以便处理多维数组/张量。

使用：

>>> import heapq
>>> import numpy
>>> a = numpy.array([1, 3, 2, 4, 5])
>>> heapq.nlargest(3, range(len(a)), a.take)
[4, 3, 1]

from operator import itemgetter
from heapq import nlargest
result = nlargest(N, enumerate(your_list), itemgetter(1))

def max_indices(arr, k):
    '''
    Returns the indices of the k first largest elements of arr
    (in descending order in values)
    '''
    assert k <= arr.size, 'k should be smaller or equal to the array size'
    arr_ = arr.astype(float)  # make a copy of arr
    max_idxs = []
    for _ in range(k):
        max_element = np.max(arr_)
        if np.isinf(max_element):
            break
        else:
            idx = np.where(arr_ == max_element)
        max_idxs.append(idx)
        arr_[idx] = -np.inf
    return max_idxs

---

我认为最具时间效率的方法是手动遍历数组并保持一个k大小的最小堆，正如其他人所提到的

我还提出了一种暴力手段：

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将返回第二大元素。

---

您可以记录这些元素的原始值，并根据需要恢复它们。

以下是查看最大元素及其位置的非常简单的方法。这里

axis

是域<代码>轴=0表示列方向的最大数量，

轴

=1表示2D情况下的行方向的最大数量。对于更高的维度，这取决于你们

M = np.random.random((3, 4))
print(M)
print(M.max(axis=1), M.argmax(axis=1))

---

此代码适用于numpy2D矩阵数组：

这会产生一个真正错误的n_最大矩阵索引，该索引还可以从矩阵数组中提取n_最大元素

三个答案，比较编码的简单性和速度 速度对我的需求很重要，所以我测试了这个问题的三个答案

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

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

编码方面：

NPE的回答是第二个最优雅的，并且足够快地满足我的需求

Fred Foos的答案要求对我的需求进行最多的重构，但速度最快。我同意这个答案，因为尽管它需要更多的工作，但也不算太糟糕，而且速度优势显著

off99555的回答是最优雅的，但却是最慢的

用于测试和比较的完整代码 输出速度报告 这是一个更复杂的问题