Skip to content

Sorting & Searching

np.sort() — return a sorted copy

import numpy as np

a = np.array([3, 1, 4, 1, 5, 9, 2, 6])
print(np.sort(a))           # [1 1 2 3 4 5 6 9]
print(a)                     # unchanged — sort returns a copy

.sort() (method) — sort in place

import numpy as np

a = np.array([3, 1, 4, 1, 5, 9, 2, 6])
a.sort()
print(a)                     # [1 1 2 3 4 5 6 9]  ← `a` itself sorted

Sort in descending order

NumPy doesn't have a reverse=True. Trick: sort ascending, then reverse:

import numpy as np

a = np.array([3, 1, 4, 1, 5, 9, 2, 6])
print(np.sort(a)[::-1])     # [9 6 5 4 3 2 1 1]

Or negate, sort, negate (for numeric only):

import numpy as np
a = np.array([3, 1, 4, 1, 5, 9, 2, 6])
print(-np.sort(-a))          # works for numeric arrays

Sort 2D — along an axis

import numpy as np

a = np.array([
    [3, 1, 4],
    [9, 2, 6],
    [5, 8, 7],
])

print("sort each row (axis=1):")
print(np.sort(a, axis=1))

print("\nsort each column (axis=0):")
print(np.sort(a, axis=0))

np.argsort() — indices that would sort

The most useful sort function. Returns indices, not values.

import numpy as np

a = np.array([30, 10, 50, 20, 40])
indices = np.argsort(a)
print("indices:", indices)        # [1 3 0 4 2]
print("sorted :", a[indices])     # [10 20 30 40 50]

Why this is useful — sort one array based on another:

import numpy as np

names  = np.array(["Carol", "Alice", "Bob", "Dave"])
scores = np.array([85, 92, 78, 88])

# Sort by score, descending
order = np.argsort(-scores)
print("ranking:")
for name, score in zip(names[order], scores[order]):
    print(f"  {name}: {score}")

np.argmin / np.argmax — index of min/max

import numpy as np

a = np.array([3, 1, 4, 1, 5, 9, 2, 6])
print("argmin:", a.argmin())     # 1 — index of first 1
print("argmax:", a.argmax())     # 5 — index of 9

With axis:

import numpy as np

# Each row is a student, each column a subject
scores = np.array([
    [85, 92, 78],
    [90, 88, 95],
    [70, 95, 80],
])

# Best subject for each student (column index of max)
print("each student's best subject:", scores.argmax(axis=1))

# Best student in each subject (row index of max)
print("each subject's best student:", scores.argmax(axis=0))

Top-K — partial sort with np.argpartition

To get the top-3 without sorting everything (faster for large arrays):

import numpy as np

rng = np.random.default_rng(0)
data = rng.integers(0, 1000, size=20)
print(data)

# Top-3 indices (in any order)
top_3_indices = np.argpartition(-data, 3)[:3]
print("top-3 values:", data[top_3_indices])

# Sorted top-3
sorted_top_3 = top_3_indices[np.argsort(-data[top_3_indices])]
print("sorted top-3:", data[sorted_top_3])

np.where() — find indices matching a condition

import numpy as np

a = np.array([1, 5, 3, 9, 2, 8, 4])
indices = np.where(a > 4)
print("indices:", indices)          # (array([1, 3, 5]),)  ← tuple
print("values :", a[indices])

For 2D arrays, np.where returns a pair (row indices, col indices):

import numpy as np

m = np.array([
    [1, 5, 3],
    [8, 2, 6],
    [4, 9, 7],
])
rows, cols = np.where(m > 5)
print("rows:", rows)
print("cols:", cols)
print("values:", m[rows, cols])

np.searchsorted() — binary search in a sorted array

import numpy as np

sorted_vals = np.array([10, 20, 30, 40, 50])

# Where would we insert 25 to keep it sorted?
print(np.searchsorted(sorted_vals, 25))           # 2
print(np.searchsorted(sorted_vals, [5, 25, 45]))   # [0 2 4]

Fast — O(log n). Useful for binning into ranges.

np.unique() — distinct values (sorted)

import numpy as np

a = np.array([3, 1, 4, 1, 5, 9, 2, 6, 5, 3])

print("unique:", np.unique(a))

vals, counts = np.unique(a, return_counts=True)
print("vals  :", vals)
print("counts:", counts)

# Get the most common
most_common = vals[counts.argmax()]
print("most common:", most_common)

np.isin() — membership test

import numpy as np

a = np.array([1, 2, 3, 4, 5, 6, 7, 8])
allowed = np.array([2, 4, 6, 8])

mask = np.isin(a, allowed)
print(mask)
print(a[mask])
print(a[~mask])      # the complement

Sorting structured / 2D data — lexsort

When you have multiple "columns" to sort by:

import numpy as np

names  = np.array(["Carol", "Alice", "Bob", "Alice", "Bob"])
ages   = np.array([30,      25,      28,    32,      27])

# Primary sort = name (ascending). Secondary sort = age (ascending).
# Note: lexsort uses the LAST key as primary.
order = np.lexsort((ages, names))

for i in order:
    print(f"  {names[i]:6}  {ages[i]}")

Mini-project — top-3 students by total score

import numpy as np

names = np.array(["Alice", "Bob", "Carol", "Dave", "Eve", "Frank"])
math    = np.array([85, 78, 92, 65, 88, 73])
science = np.array([90, 82, 88, 70, 92, 80])
english = np.array([78, 85, 80, 88, 75, 90])

totals = math + science + english
print("Totals:", totals)

# Top 3 indices
top3 = np.argsort(-totals)[:3]
print("\nTop 3:")
for rank, i in enumerate(top3, 1):
    print(f"  {rank}. {names[i]:6}{totals[i]}")

Cheatsheet

Want Use
Sorted copy np.sort(arr)
Sort in place arr.sort()
Sorted indices np.argsort(arr)
Descending np.sort(arr)[::-1]
Top-K (fast, unordered) np.argpartition(-arr, K)[:K]
Index of min/max arr.argmin() / arr.argmax()
Find condition matches np.where(arr > 5)
Binary search in sorted array np.searchsorted(sorted, vals)
Unique values + counts np.unique(arr, return_counts=True)
Is in a set np.isin(arr, allowed)
Multi-key sort np.lexsort((second_key, primary_key))

Common pitfalls

  • np.sort() returns a copy; .sort() is in place — easy mix-up.
  • np.where returns a tuple — even for 1D. Use np.where(cond)[0] if you need a plain array of indices.
  • Sort doesn't have a reverse=True — use [::-1] or negate.
  • argsort on strings sorts lexicographically"10" comes before "2". Convert to numbers first.
  • lexsort order surprise — last key is the primary key, not the first.

Practice

What does this print?

Expected: [1 3 0 4 2]

import numpy as np
a = np.array([30, 10, 50, 20, 40])
print(np.argsort(a))

Sort by score, descending (highest first)

Expected: [95 92 88 85]

import numpy as np
scores = np.array([85, 92, 88, 95])
print(np.sort(scores))           # bug: sorts ascending — need to reverse

Quiz — Quick check

What you remember

Q1. Difference between np.sort(a) and a.sort()?

  • np.sort returns a sorted copy; a.sort() sorts in place
  • They're identical
  • np.sort is for 2D only
  • a.sort() returns the sorted array

Why: a.sort() modifies a and returns None. np.sort(a) returns a new sorted array, leaving a unchanged.

Q2. What does np.argsort(arr) return?

  • The sorted values
  • The indices that would sort arr
  • The minimum value
  • A boolean mask

Why: argsort gives you the order. arr[np.argsort(arr)] is equivalent to np.sort(arr). It's powerful because you can apply the same order to other arrays (like sorting names by score).

Q3. How do you sort in descending order?

  • np.sort(arr, reverse=True)
  • np.sort(arr)[::-1] or -np.sort(-arr) (for numeric)
  • np.argsort(arr)
  • arr.sort(desc=True)

Why: NumPy's sort doesn't have a reverse argument. The idioms are slicing the result with [::-1] or negating numeric arrays before sort.

Common doubts

When should I use argpartition instead of argsort?

For top-K when K is much smaller than N. argsort sorts the entire array (O(n log n)). argpartition only ensures the top-K elements are in the first K positions (O(n)) — they're not sorted relative to each other. Use argpartition for "top 10 of 1 million," then argsort the smaller result if you need ranking.

Why does np.where(cond) return a tuple?

Because np.where always returns one array per dimension. For 1D, it's a 1-tuple: (array([1, 3, 5]),). For 2D, it's (row_indices, col_indices). If you want a plain array of indices for 1D, use np.where(cond)[0].

How is np.searchsorted different from np.where?

np.where does a linear scan of a boolean array — O(n). np.searchsorted does a binary search on a sorted array — O(log n). Use searchsorted when you need fast lookups against a sorted reference (like binning, percentile lookups, etc.).

What's next

Linear Algebra