HowToCodePage

Sorting, searching and set operations

Sort along the axis you mean, find insertion points without a loop, and use NumPy's set routines for joins and membership.

sort, argsort and partial sorts

np.sort returns a sorted copy; arr.sort() sorts in place. When you need the permutation rather than the values, use np.argsort — it is the bridge between sorting and indexing.

a = np.array([3, 1, 2])
np.sort(a)                 # array([1, 2, 3]) - a is unchanged
a.sort()                   # in place

idx = np.argsort(a)        # the indices that would sort a
a[idx]                     # same as np.sort(a)

m = np.array([[3, 1], [2, 9]])
np.sort(m, axis=0)         # sort each column
np.sort(m, axis=1)         # sort each row
np.sort(m, axis=None)      # flatten first, then sort
np.sort(a)[::-1]           # descending

np.partition(a, 2)         # k smallest on the left, rest on the right
top3 = a[np.argpartition(a, -3)[-3:]]   # O(n), unsorted top three
FunctionCostUse for
np.sort / argsortO(n log n)Full ordering
np.partition / argpartitionO(n)A handful of extremes
np.lexsortO(n log n)Sort by several keys, last key primary
np.sort(a, kind="stable")O(n log n)Preserving the order of equal keys

np.searchsorted is a binary search: give it a sorted array and a value, it returns the index where the value belongs. With an array of queries it answers thousands of lookups at once.

grades = np.array([10, 20, 30, 40])   # must be sorted

np.searchsorted(grades, 25)                    # 2 - insert position
np.searchsorted(grades, [5, 35, 40], side="right")   # array([0, 3, 4])

# bucket thousands of scores into grade bands in one call
scores = np.array([7, 22, 41, 15])
band = np.searchsorted(grades, scores, side="right")   # array([0, 1, 3, 1])

np.nonzero(a > 1)          # tuple of index arrays, one per dimension
np.where(a > 1)            # same thing for a single condition

Set operations

x = np.array([3, 1, 1, 2])

vals, counts = np.unique(x, return_counts=True)
# vals = [1, 2, 3], counts = [2, 1, 1]
np.unique(x, return_index=True)     # first occurrence of each value

np.isin(np.array([1, 5, 9]), np.array([1, 2, 3]))   # array([True, False, False])
np.intersect1d(x, np.array([2, 3]))   # array([2, 3])
np.setdiff1d(x, np.array([2]))        # array([1, 3])
np.union1d(x, np.array([9]))          # array([1, 2, 3, 9])

# membership without a Python loop
mask = np.isin(big_ids, allowed_ids)
big[mask]
💡
np.unique sorts internally, so it is O(n log n) and returns values in ascending order. That ordering is a feature — but for millions of repeated categories a hash-based tool such as pandas is faster.

FAQ

How do I get the n largest values?
a[np.argpartition(a, -n)[-n:]] is O(n) and beats a full sort. Sort the small result afterwards if you want it ordered.
Why does searchsorted give the wrong index?
The first argument must already be sorted, and the side argument decides whether equal values land before or after existing ones. Both are easy to get subtly wrong.

Debugging array code: shape errors and NumPy 2 pitfalls Performance: einsum, ufunc tricks and avoiding copies

Last refreshed 2026-09-18.