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.wherereturns a tuple — even for 1D. Usenp.where(cond)[0]if you need a plain array of indices. - ❗ Sort doesn't have a
reverse=True— use[::-1]or negate. - ❗
argsorton strings sorts lexicographically —"10"comes before"2". Convert to numbers first. - ❗
lexsortorder surprise — last key is the primary key, not the first.
Practice¶
What does this print?
Expected: [1 3 0 4 2]
Sort by score, descending (highest first)
Expected: [95 92 88 85]
Quiz — Quick check¶
What you remember
Q1. Difference between np.sort(a) and a.sort()?
-
np.sortreturns a sorted copy;a.sort()sorts in place - They're identical
-
np.sortis for 2D only -
a.sort()returns the sorted array
Why:
a.sort()modifiesaand returnsNone.np.sort(a)returns a new sorted array, leavingaunchanged.
Q2. What does np.argsort(arr) return?
- The sorted values
- The indices that would sort
arr - The minimum value
- A boolean mask
Why:
argsortgives you the order.arr[np.argsort(arr)]is equivalent tonp.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
reverseargument. 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.).