如何获取 NumPy 数组中 N 个最大值的索引？-IT科技

摘要：问题描述：NumPy 提出了一种通过获取数组最大值的索引的方法np.argmax。我想要类似的东西，但返回N最大值的索引。例如，如果我有一个数组[1, 3, 2, 4, 5]，那么将返回与元素相对应的nargmax(array, n=3)索引。[4, 3, 1]`[5, 4, 3]`解决方案 1：较新的 Nu...

问题描述：

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

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

例如，如果我有一个数组[1, 3, 2, 4, 5]，那么将返回与元素相对应的nargmax(array, n=3)索引。[4, 3, 1]`[5, 4, 3]`

解决方案 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 ) 时间。

解决方案 2：

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

>>> import numpy as np
>>> arr = np.array([1, 3, 2, 4, 5])
>>> arr.argsort()[-3:][::-1]
array([4, 3, 1])

这涉及对数组进行完整排序。我想知道是否numpy提供了内置方法来进行部分排序；到目前为止我还没有找到。

如果这个解决方案太慢（特别是对于小型的n），可能值得研究在Cython中编写一些代码。

解决方案 3：

更简单：

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

其中n是最大值的数量。

解决方案 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——堆队列算法

解决方案 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])

解决方案 6：

三种答案在编码难易度和速度方面的比较

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

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

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

编码方式：

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的回答：

[(631934, 999999), (788104, 999998), (413003, 999997), (536514, 999996), (81029, 999995)]
0.1349949836730957

Fred Foo 的回答：

[(631934, 999999), (788104, 999998), (413003, 999997), (536514, 999996), (81029, 999995)]
0.011161565780639648

off99555的回答：

[(631934, 999999), (788104, 999998), (413003, 999997), (536514, 999996), (81029, 999995)]
0.439760684967041

解决方案 7：

如果您不关心第 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

解决方案 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]])

解决方案 9：

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

在此处输入图片描述

我们可以看到，如果您想要严格升序排列前 k 个索引，np.argpartition则不会返回您想要的内容。

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

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

在此处输入图片描述

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

解决方案 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])

当然，这涉及到篡改您的原始数组。您可以通过复制或替换原始值来修复（如果需要）。...无论哪种方式对您的用例来说更便宜。

解决方案 11：

使用：

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

现在result列表将包含N元组 ( index, value)，其中value最大化。

解决方案 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])

解决方案 13：

使用argpartition的矢量化二维实现：

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

解决方案 14：

我发现它使用起来非常直观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]

解决方案 15：

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

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

解决方案 16：

这是一种更复杂的方法，如果第 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])

解决方案 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 将返回第二大元素。您可以记录这些元素的原始值，并在需要时恢复它们。

解决方案 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 元素

解决方案 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)

解决方案 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]

解决方案 21：

如果您正在处理 NaN 和/或在理解 np.argpartition 方面有问题，请尝试pandas.DataFrame.sort_values。

import numpy as np
import pandas as pd    

a = np.array([9, 4, 4, 3, 3, 9, 0, 4, 6, 0])

df = pd.DataFrame(a, columns=['array'])
max_values = df['array'].sort_values(ascending=False, na_position='last')
ind = max_values[0:3].index.to_list()

此示例给出了 3 个最大非 NaN 值的索引。可能效率不高，但易于阅读和自定义。