How Can I Efficiently Find the Indices of Multiple Maximum Values in a NumPy Array?

Retrieving Indices of Multiple Maximum Values in NumPy Arrays

NumPy provides a convenient np.argmax function for retrieving the index of the maximum value in an array. However, what if you need to find the indices of the top N maximum values?

Solution

Recent NumPy versions (1.8 and above) introduce the argpartition function for this purpose. To obtain the indices of the top N elements, follow these steps:

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)

Explanation

np.argpartition sorts the array partially, partitioning it into two sub-arrays: the first sub-array contains the top N elements (in this case, the largest 4 elements), and the second sub-array contains the remaining elements. The returned array ind contains the indices of the elements in the first sub-array.

The output in this example would be:

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

Optimizations

If sorted indices are also needed, you can sort them separately:

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

This step requires O(k log k) time, where k is the number of top elements to retrieve. Overall, this approach has a time complexity of O(n k log k), making it efficient for large arrays and moderate values of k.

The above is the detailed content of How Can I Efficiently Find the Indices of Multiple Maximum Values in a NumPy Array?. For more information, please follow other related articles on the PHP Chinese website!