Home > Backend Development > Python Tutorial > How Can I Efficiently Find the Indices of Multiple Maximum Values in a NumPy Array?

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

Patricia Arquette
Release: 2024-12-02 13:24:15
Original
901 people have browsed it

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)
Copy after login

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]
Copy after login

Optimizations

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

sorted_ind = ind[np.argsort(a[ind])]
Copy after login

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!

source:php.cn
Statement of this Website
The content of this article is voluntarily contributed by netizens, and the copyright belongs to the original author. This site does not assume corresponding legal responsibility. If you find any content suspected of plagiarism or infringement, please contact admin@php.cn
Latest Articles by Author
Popular Tutorials
More>
Latest Downloads
More>
Web Effects
Website Source Code
Website Materials
Front End Template