How to Find the Indices of Multiple Maximum Values in a NumPy Array?

Retrieving Indices of Multiple Maximum Values in a NumPy Array

NumPy arrays provide the np.argmax function to locate the index of the maximum element. However, if you require the indices of N maximum values, consider the following solutions:

Recent NumPy Versions:

For NumPy versions 1.8 and above, the argpartition function offers an efficient method:

import numpy as np

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

# Get top 4 indices
n_max = 4
ind = np.argpartition(a, -n_max)[-n_max:]

# Retrieve top 4 values
top_max = a[ind]

Older NumPy Versions:

Prior to NumPy 1.8, you can use the argsort function as follows:

# Get top 4 indices
n_max = 4
ind = np.argsort(a)[-n_max:]

# Retrieve top 4 values
top_max = a[ind]

Sorting the Indices:

By default, argpartition returns unsorted indices. If you require sorted indices, use:

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

Time Complexity:

• argpartition: O(n) in the worst case
• argsort: O(n log n)
• Combined approach (sorted indices): O(n k log k) for top-k sorted elements

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