Sorting and Searching

NumPy's sorting and searching functions run in C with tuned algorithms - quicksort by default, with mergesort (stable) and heapsort as alternatives - and the search family (searchsorted, unique) builds directly on sorted arrays. The patterns matter more than the functions: sorting indices with argsort to reorder several arrays consistently, argpartition for top-k without a full sort, and searchsorted as the vectorized binary search that turns sorted tables into O(log n) lookups.

Sorting Arrays

import numpy as np

a = np.array([3, 1, 4, 1, 5, 9, 2, 6])
print(np.sort(a))          # new array, sorted - a is untouched
a.sort()                   # in-place, sorts a itself
print(a)

# 2-D: sort each row (axis=1) or each column (axis=0)
m = np.array([[3, 1], [2, 4], [5, 0]])
print(np.sort(m, axis=1))  # sort within each row
print(np.sort(m, axis=0))  # sort within each column

np.sort returns a new array; a.sort() sorts in place and returns None (the list convention). For N-D arrays, axis names the axis to sort along - the other axes keep their positions, and each 1-D slice along that axis is sorted independently. The default algorithm is introsort (a quicksort/heapsort hybrid, not stable); pass kind="mergesort" when equal elements must keep their original order.

Sorting by Permutation: argsort

scores = np.array([0.8, 0.3, 0.95, 0.1, 0.7])
names = np.array(["ada", "bo", "cy", "de", "el"])

order = np.argsort(scores)            # indices that sort scores
print(order)                           # [3 1 4 0 2] - ascending by score
print(scores[order])                   # [0.1 0.3 0.7 0.8 0.95]

# reorder EVERY array by the same permutation - the aligned-sort idiom
print(names[order])                    # names in the same score order
print(names[np.argsort(scores)[::-1]]) # best score first

argsort returns the permutation of indices that would sort the array. That permutation is the tool: apply it to any number of parallel arrays (names, dates, feature rows) and everything stays aligned with the sorted order. This one idiom replaces what would be a decorate-sort-undecorate dance in plain Python, and it scales because the sort itself is in C.

Top-k Without Sorting Everything: argpartition

a = np.random.default_rng(0).normal(size=1_000_000)
k = 10

# full sort: O(n log n) - overkill when you only need the top k
# partition: O(n) - the k largest land in the last k slots (unordered)
idx = np.argpartition(a, -k)[-k:]
print(a[idx].sort())          # the k largest values, now ordered
print(len(idx))               # 10

argpartition(a, kth) rearranges so that the element at kth is where it would be in sorted order, with everything smaller before it and larger after - no other ordering guaranteed. Slicing the last k positions gives the k largest (or first k for the smallest) in a single O(n) pass. For "top 10 of a million" this is the standard optimization, and the final small sort of the k elements is trivial.

Binary Search on Sorted Data

table = np.array([10, 20, 30, 40, 50])
queries = np.array([15, 20, 35, 55])

pos = np.searchsorted(table, queries)      # first index where each fits
print(pos)                                  # [1 1 3 5]

# the lookup idiom: banded / threshold mapping
def band(x):
    return np.searchsorted([0, 50, 80], x)  # 0, 1, or 2 per value
print(band(queries))                        # [0 0 1 2]

# membership on a sorted array
print(np.searchsorted(table, 30) < len(table)
      and table[np.searchsorted(table, 30)] == 30)

searchsorted is a vectorized binary search: for each query value it returns the insertion index in the sorted table, all in one C call. That single function covers threshold mapping (which band does this value fall in), nearest-value lookups (index then step left or right), and membership tests. The precondition is a sorted table - argsort first, then searchsorted forever after, is the standard setup.

Unique Values and Counts

a = np.array([3, 1, 3, 2, 1, 3, 4])
print(np.unique(a))                    # [1 2 3 4] - sorted unique values

u, inverse = np.unique(a, return_inverse=True)
print(inverse)                         # maps each element to its slot in u
print(u[inverse] == a).all()           # True - the round trip holds

u2, counts = np.unique(a, return_counts=True)
print(dict(zip(u2.tolist(), counts.tolist())))   # {1:2, 2:1, 3:3, 4:1}

unique returns the sorted set of distinct values. The two return flags are the power: return_inverse gives the index array that reconstructs the original (the encode/decode pattern - turn labels into small integers, process, decode back), and return_counts is the histogram for one value set, the standard first step of categorical analysis.

Summary

Sort with np.sort (new array) or a.sort() (in place), and argsort for the permutation that reorders parallel arrays in lockstep. Use argpartition for top-k in linear time instead of sorting a million elements. Keep lookup tables sorted and answer lookups with the vectorized searchsorted - thresholds, bands, nearest values. And use unique with its inverse/counts flags for the encode-decode and frequency patterns that categorical data always produces. All of it runs in C; the Python layer is just choosing the right index permutation.