numpy
python
array manipulation
maximum values
indexing

How do I get indices of N maximum values in a NumPy array?

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

Finding the indices of the top N values in a NumPy array is a common ranking task in data analysis and machine learning. The main design choice is whether you need the top indices fully sorted or whether you only care about identifying the top subset efficiently.

NumPy gives you two main tools for this. argsort is simpler and returns a fully ordered result. argpartition is usually faster for large arrays when you only need the top few items.

Use argsort for the Clearest Answer

If the array is not huge or clarity matters more than micro-optimization, use argsort:

python
1import numpy as np
2
3a = np.array([4, 1, 9, 7, 9, 3])
4n = 3
5
6idx = np.argsort(a)[-n:][::-1]
7print(idx)
8print(a[idx])

This sorts the full array of indices, takes the last n, and reverses them so the largest values come first.

The output indices are ranked, which is often exactly what you want in reporting or post-processing code.

Use argpartition for Large Arrays

If the array is large and you only care about the top N, argpartition avoids a full sort:

python
1import numpy as np
2
3a = np.array([4, 1, 9, 7, 9, 3])
4n = 3
5
6idx = np.argpartition(a, -n)[-n:]
7print(idx)
8print(a[idx])

This finds the correct top subset, but the subset is not guaranteed to be internally sorted. If you need the subset ranked afterward, sort only that smaller portion:

python
idx = idx[np.argsort(a[idx])[::-1]]
print(idx)
print(a[idx])

That is often the best balance between speed and clear output.

Ties Need a Policy

If the array contains repeated maximum values, NumPy will return valid indices, but tied elements may not appear in the order you expect. If deterministic ordering matters, define a tie-breaker explicitly.

A simple pattern is to sort by value descending and index ascending within the selected subset:

python
1a = np.array([9, 4, 9, 7, 9, 3])
2n = 3
3
4idx = np.argpartition(a, -n)[-n:]
5idx = idx[np.lexsort((idx, -a[idx]))]
6
7print(idx)
8print(a[idx])

This makes the output stable instead of depending on internal partition details.

Multi-Dimensional Arrays

For a 2D or higher-dimensional array, decide whether you want the top N globally or along an axis.

Global top N example:

python
1m = np.array([
2    [1, 8, 3],
3    [7, 2, 9],
4    [4, 6, 5],
5])
6
7n = 4
8flat_idx = np.argpartition(m.ravel(), -n)[-n:]
9rows, cols = np.unravel_index(flat_idx, m.shape)
10values = m[rows, cols]
11
12order = np.argsort(values)[::-1]
13print(list(zip(rows[order], cols[order], values[order])))

That is different from finding the top N in each row or column, so make the requirement explicit before choosing the method.

Handle Edge Cases Deliberately

A reusable helper should decide what happens when:

  • 'n is zero'
  • 'n is negative'
  • 'n is larger than the array size'
  • the array is empty

Example:

python
1def top_n_indices(a: np.ndarray, n: int) -> np.ndarray:
2    if n <= 0 or a.size == 0:
3        return np.array([], dtype=int)
4    n = min(n, a.size)
5    idx = np.argpartition(a, -n)[-n:]
6    return idx[np.argsort(a[idx])[::-1]]

This keeps behavior predictable instead of relying on accidental slicing behavior.

Common Pitfalls

The biggest mistake is assuming argpartition returns the top indices in sorted order. It does not. It only guarantees that the chosen subset contains the right elements.

Another common issue is using a full argsort on very large arrays when only a small top subset is needed.

Developers also forget that ties may need explicit handling if the output order must be stable or reproducible.

Finally, for multi-dimensional arrays, be clear about whether "top N" means globally or along an axis. Those are different operations.

Summary

  • Use np.argsort when you want the simplest fully ranked top-N solution.
  • Use np.argpartition when performance matters and you only need the top subset.
  • Sort the selected subset afterward if ranked output is required.
  • Define a tie-breaking rule when deterministic ordering matters.
  • Be explicit about axis behavior and edge cases in reusable helpers.

Course illustration
Course illustration

All Rights Reserved.