python 如何在numpy數組中獲取N個最大值的索引?
較新的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 )時間。
Numpy提出了一種通過np.argmax
數組最大值索引的np.argmax
。
我想要一個類似的東西,但返回N個最大值的索引。
例如,如果我有一個數組[1, 3, 2, 4, 5]
,它的function(array, n=3)
將返回[4, 3, 1]
。
謝謝 :)
更簡單:
idx = (-arr).argsort()[:n]
其中n是最大值的數量。
如果你不關心第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值的增加。
from operator import itemgetter
from heapq import nlargest
result = nlargest(N, enumerate(your_list), itemgetter(1))
現在結果列表將包含N個元組(索引,值),其中值被最大化
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的下一次調用將返回第二大元素。 如果需要,您可以記錄這些元素的原始值並恢復它們。