如何在 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