[python] 如何在numpy數組中獲取N個最大值的索引?



Answers

較新的NumPy版本(1.8及以上版本)有一個稱為argpartition的功能。 要獲得四個最大元素的指標,請執行

>>> 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不同的argsort ,這個函數在最壞的情況下以線性時間運行,但是返回的索引沒有被排序,從評估a[ind]的結果可以看出。 如果您也需要,請在之後進行分類:

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

以這種方式獲得排序順序的top- k元素需要O( n + k log k )時間。

Question

Numpy提出了一種通過np.argmax數組最大值索引的np.argmax

我想要一個類似的東西,但返回N個最大值的索引。

例如,如果我有一個數組[1, 3, 2, 4, 5] ,它的function(array, n=3)將返回[4, 3, 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

也適用於二維數組。 例如

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



正如其他人所提到的,我認為時間效率最高的方式是通過數組手動迭代並保持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的下一次調用將返回第二大元素。 如果需要,您可以記錄這些元素的原始值並恢復它們。




如果你不關心第K個最大元素的順序 ,你可以使用argpartition ,它應該比通過argsort完全排序argsort

K = 4 # we want the indeces 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值的增加。




更簡單:

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

其中n是最大值的數量。




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

現在結果列表將包含N個元組(索引,值),其中值被最大化




bottleneck具有部分排序功能,如果為了獲得N個最大值而排序整個數組的代價太大。

我對這個模塊一無所知。 我剛剛GOOGLE了numpy partial sort




Links



Tags

python python   numpy