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

Question

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

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

例如，如果我有一个数组[1, 3, 2, 4, 5] ，那么nargmax(array, n=3)将返回与元素[5, 4, 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])

>>> 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 ) 时间。

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

学分去这个问题。

我进行了一些测试，随着数组大小和 K 值的增加，看起来argpartition优于argsort 。

Answer 7

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

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

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

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

编码明智：

NPE 的答案是我需要的下一个最优雅和足够快的答案。
Fred Foos 的回答需要对我的需求进行最多的重构，但速度最快。 我接受了这个答案，因为即使它需要更多的工作，它也不算太糟糕并且具有显着的速度优势。
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

方法np.argpartition只返回 k 个最大的索引，执行本地排序，并且在数组很大时比np.argsort （执行完整排序）更快。 但返回的索引不是升序/降序。 让我们举个例子：

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

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

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

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

Answer 10

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

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

利用：

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

现在result列表将包含N个元组（ index ， value ），其中value被最大化。

Answer 12

利用：

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

它也适用于二维数组。 例如，

In [0]: A = np.array([[ 0.51845014,  0.72528114],
                     [ 0.88421561,  0.18798661],
                     [ 0.89832036,  0.19448609],
                     [ 0.89832036,  0.19448609]])
In [1]: max_indices(A, 8)
Out[1]:
    [(array([2, 3], dtype=int64), array([0, 0], dtype=int64)),
     (array([1], dtype=int64), array([0], dtype=int64)),
     (array([0], dtype=int64), array([1], dtype=int64)),
     (array([0], dtype=int64), array([0], dtype=int64)),
     (array([2, 3], dtype=int64), array([1, 1], dtype=int64)),
     (array([1], dtype=int64), array([1], dtype=int64))]

In [2]: A[max_indices(A, 8)[0]][0]
Out[2]: array([ 0.89832036])

Answer 13

以下是查看最大元素及其位置的一种非常简单的方法。 这里axis是域； axis = 0 表示按列的最大数量， axis = 1 表示 2D 情况下的按行的最大数量。 而对于更高的维度，这取决于你。

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

Answer 14

这是一种更复杂的方法，如果第 n 个值有关系，则增加 n：

>>>> def get_top_n_plus_ties(arr,n):
>>>>     sorted_args = np.argsort(-arr)
>>>>     thresh = arr[sorted_args[n]]
>>>>     n_ = np.sum(arr >= thresh)
>>>>     return sorted_args[:n_]
>>>> get_top_n_plus_ties(np.array([2,9,8,3,0,2,8,3,1,9,5]),3)
array([1, 9, 2, 6])

Answer 15

bottleneck有部分排序功能，如果只是为了得到 N 个最大值而对整个数组进行排序的代价太大了。

_{我对这个模块一无所知；} _{我只是在谷numpy partial sort搜索了numpy partial sort 。}

Answer 16

我发现使用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]

Answer 17

正如其他人所提到的，我认为最省时的方法是手动遍历数组并保持 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 的下一次调用将返回第二大元素。 如果需要，您可以记录这些元素的原始值并恢复它们。

Answer 18

此代码适用于 numpy 2D 矩阵数组：

mat = np.array([[1, 3], [2, 5]]) # numpy matrix
 
n = 2  # n
n_largest_mat = np.sort(mat, axis=None)[-n:] # n_largest 
tf_n_largest = np.zeros((2,2), dtype=bool) # all false matrix
for x in n_largest_mat: 
  tf_n_largest = (tf_n_largest) | (mat == x) # true-false  

n_largest_elems = mat[tf_n_largest] # true-false indexing

这会产生一个真假 n_largest 矩阵索引，它也可以从矩阵数组中提取 n_largest 元素

Answer 19

当 top_k<<axis_length 时，它优于 argsort。

import numpy as np

def get_sorted_top_k(array, top_k=1, axis=-1, reverse=False):
    if reverse:
        axis_length = array.shape[axis]
        partition_index = np.take(np.argpartition(array, kth=-top_k, axis=axis),
                                  range(axis_length - top_k, axis_length), axis)
    else:
        partition_index = np.take(np.argpartition(array, kth=top_k, axis=axis), range(0, top_k), axis)
    top_scores = np.take_along_axis(array, partition_index, axis)
    # resort partition
    sorted_index = np.argsort(top_scores, axis=axis)
    if reverse:
        sorted_index = np.flip(sorted_index, axis=axis)
    top_sorted_scores = np.take_along_axis(top_scores, sorted_index, axis)
    top_sorted_indexes = np.take_along_axis(partition_index, sorted_index, axis)
    return top_sorted_scores, top_sorted_indexes

if __name__ == "__main__":
    import time
    from sklearn.metrics.pairwise import cosine_similarity

    x = np.random.rand(10, 128)
    y = np.random.rand(1000000, 128)
    z = cosine_similarity(x, y)
    start_time = time.time()
    sorted_index_1 = get_sorted_top_k(z, top_k=3, axis=1, reverse=True)[1]
    print(time.time() - start_time)

Answer 20

您可以简单地使用字典来查找 numpy 数组中的前 k 个值和索引。 例如，如果您想查找前 2 个最大值和索引

import numpy as np
nums = np.array([0.2, 0.3, 0.25, 0.15, 0.1])


def TopK(x, k):
    a = dict([(i, j) for i, j in enumerate(x)])
    sorted_a = dict(sorted(a.items(), key = lambda kv:kv[1], reverse=True))
    indices = list(sorted_a.keys())[:k]
    values = list(sorted_a.values())[:k]
    return (indices, values)

print(f"Indices: {TopK(nums, k = 2)[0]}")
print(f"Values: {TopK(nums, k = 2)[1]}")


Indices: [1, 2]
Values: [0.3, 0.25]

Answer 21

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

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

问题描述

20 个解决方案

解决方案1
861 已采纳

解决方案2
482 2011-08-02 10:32:53

解决方案3
71 2014-12-11 22:13:05

解决方案4
47 2013-09-09 05:30:32

解决方案5
42 2016-08-10 21:42:27

解决方案6
13 2016-05-13 13:16:28

解决方案7
12 2020-11-07 05:35:49

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

完整的测试和比较代码

输出速度报告

解决方案8
10 2016-12-11 14:34:34

解决方案9
7 2018-01-25 05:00:29

解决方案10
5 2011-08-02 13:54:12

解决方案11
4 2016-04-17 10:06:35

解决方案12
4 2018-01-30 14:15:38

解决方案13
1 2018-06-16 08:20:27

解决方案14
1 2020-11-19 20:57:25

解决方案15
1 2011-08-02 10:37:08

解决方案16
0 2018-01-12 14:39:40

解决方案17
0 2018-04-25 10:09:07

解决方案18
0 2019-10-23 04:28:40

解决方案19
0 2021-01-13 09:05:10

解决方案20
0 2021-08-25 19:15:07

解决方案21
0 2022-01-27 19:48:54

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

问题描述

20 个解决方案

解决方案1 861 已采纳

解决方案2 482 2011-08-02 10:32:53

解决方案3 71 2014-12-11 22:13:05

解决方案4 47 2013-09-09 05:30:32

解决方案5 42 2016-08-10 21:42:27

解决方案6 13 2016-05-13 13:16:28

解决方案7 12 2020-11-07 05:35:49

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

完整的测试和比较代码

输出速度报告

解决方案8 10 2016-12-11 14:34:34

解决方案9 7 2018-01-25 05:00:29

解决方案10 5 2011-08-02 13:54:12

解决方案11 4 2016-04-17 10:06:35

解决方案12 4 2018-01-30 14:15:38

解决方案13 1 2018-06-16 08:20:27

解决方案14 1 2020-11-19 20:57:25

解决方案15 1 2011-08-02 10:37:08

解决方案16 0 2018-01-12 14:39:40

解决方案17 0 2018-04-25 10:09:07

解决方案18 0 2019-10-23 04:28:40

解决方案19 0 2021-01-13 09:05:10

解决方案20 0 2021-08-25 19:15:07

解决方案21 0 2022-01-27 19:48:54

解决方案1
861 已采纳

解决方案2
482 2011-08-02 10:32:53

解决方案3
71 2014-12-11 22:13:05

解决方案4
47 2013-09-09 05:30:32

解决方案5
42 2016-08-10 21:42:27

解决方案6
13 2016-05-13 13:16:28

解决方案7
12 2020-11-07 05:35:49

解决方案8
10 2016-12-11 14:34:34

解决方案9
7 2018-01-25 05:00:29

解决方案10
5 2011-08-02 13:54:12

解决方案11
4 2016-04-17 10:06:35

解决方案12
4 2018-01-30 14:15:38

解决方案13
1 2018-06-16 08:20:27

解决方案14
1 2020-11-19 20:57:25

解决方案15
1 2011-08-02 10:37:08

解决方案16
0 2018-01-12 14:39:40

解决方案17
0 2018-04-25 10:09:07

解决方案18
0 2019-10-23 04:28:40

解决方案19
0 2021-01-13 09:05:10

解决方案20
0 2021-08-25 19:15:07

解决方案21
0 2022-01-27 19:48:54