如何有效率地尋找 NumPy 陣列中多個最大值的索引？

檢索NumPy 數組中多個最大值的索引

NumPy 提供了一個方便的np.argmax 函數來檢索NumPy 數組中最大值的索引一個數組。但是，如果您需要找到前 N 個最大值的索引怎麼辦？

解決方案

最近的 NumPy 版本（1.8 及更高版本）為此引入了 argpartition 函數目的。若要取得前N 個元素的索引，請依照下列步驟操作：

import numpy as np

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

# Find indices of top N elements (N = 4 in this case)
ind = np.argpartition(a, -4)[-4:]

# Extract top N elements
top4 = a[ind]

# Print indices and top N elements
print("Indices:", ind)
print("Top 4 elements:", top4)

說明

np.argpartition 將陣列進行部分排序，將其分為兩個子數組：第一個子數組包含前N 個元素（在本例中為最大的4 個元素），第二個子數組包含其餘元素。傳回的數組 ind 包含第一個子數組中元素的索引。

此範例中的輸出將是：

Indices: [1 5 8 0]
Top 4 elements: [4 9 6 9]

最佳化

如果還需要排序索引，可以單獨排序：

sorted_ind = ind[np.argsort(a[ind])]

這一步驟需要O(k log k) 時間，其中k 是要檢索的頂部元素的數量。總的來說，這種方法的時間複雜度為 O(n k log k)，對於大型陣列和中等 k 值非常有效。

以上是如何有效率地尋找 NumPy 陣列中多個最大值的索引？的詳細內容。更多資訊請關注PHP中文網其他相關文章！