NumPy Cheatsheet

Useful Functions

Use this NumPy reference while you build software engineering projects, review code, or refresh the syntax you reach for most.

Sorting and Searching

import numpy as np

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

# Sorting — returns NEW sorted array
np.sort(a)                        # ascending
np.sort(a)[::-1]                  # descending (reverse after sort)
a.sort()                           # in-place sort (no return value)

# Multi-axis
b = np.array([[3, 1], [4, 2]])
np.sort(b, axis=0)                # sort each column
np.sort(b, axis=1)                # sort each row
np.sort(b, axis=None)             # flatten then sort

# Stable sort (preserves relative order of equal elements)
np.sort(a, kind="stable")         # mergesort is stable
np.sort(a, kind="mergesort")      # explicit
np.sort(a, kind="quicksort")      # default (not stable but fast)
np.sort(a, kind="heapsort")

# Indirect sort — indices that would sort the array
np.argsort(a)                     # default kind="quicksort" (introsort) — NOT stable
np.argsort(a, kind="stable")      # stable — required for equal-element order
a[np.argsort(a)]                  # same as np.sort(a)

# Partial sort — nth element in its sorted position, no guarantees otherwise
np.partition(a, 3)                # element at index 3 is in its final place
np.argpartition(a, 3)             # indices version

# Lexicographic sort (multiple keys)
names  = np.array(["Bob", "Alice", "Charlie", "Alice"])
scores = np.array([90, 85, 95, 92])
idx = np.lexsort((scores, names))   # sort by names first, then scores
# keys are specified right-to-left (last key = primary sort key)
names[idx]    # sorted result

Searching

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

np.searchsorted(np.sort(a), 4)             # bisect left: 3   (sorted: [1,1,3,4,5,9])
np.searchsorted(np.sort(a), 4, side="right")  # bisect right: 4
np.searchsorted([1, 3, 5, 7], [2, 4, 6])  # [1, 2, 3]

np.where(a > 3)          # (array([2, 4, 5]),)  — indices where True
np.flatnonzero(a > 3)    # [2, 4, 5]           — flat indices
np.argmax(a)             # 5
np.argmin(a)             # 1

Set Operations

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

np.unique(a)                     # [1, 2, 3]
np.unique(a, return_counts=True) # ([1,2,3], [2,2,1])
np.intersect1d(a, b)             # [2, 3]
np.union1d(a, b)                 # [1, 2, 3, 4, 5]
np.setdiff1d(a, b)               # [1]    in a but not b
np.setxor1d(a, b)                # [1, 4, 5]  in one but not both
np.in1d(a, b)                    # [F, T, T, T, F]  deprecated alias
np.isin(a, b)                    # [F, T, T, T, F]  preferred
np.isin(a, b, invert=True)       # complement

Type Testing and Conversion

np.isscalar(3.0)          # True
np.isscalar(np.array(3))  # False
np.iterable(a)            # True

np.isreal(a)              # True for real-valued elements
np.iscomplex(a)           # False
np.isrealobj(a)           # True — array has no imaginary part
np.iscomplexobj(a)        # False

np.can_cast("int32", "float64")       # True
np.can_cast("float64", "int32")       # False (would lose info)
np.can_cast("float64", "int32", casting="unsafe")  # True

np.result_type(np.float32, np.int64)   # float64
np.common_type(np.array([1.0]), np.array([1+0j]))  # complex128 — prefer np.result_type
np.min_scalar_type(1000)               # uint16 (smallest type that fits)

Copying and Memory

np.copy(a)               # explicit copy (same as a.copy())
np.asarray(a)            # no-copy if already ndarray of correct dtype/order
np.ascontiguousarray(a)  # C-contiguous, copy if needed
np.asfortranarray(a)     # F-contiguous
np.require(a, dtype=float, requirements=["C", "W"])  # enforce requirements

# requirements: "C" contiguous, "F" Fortran, "O" own data,
#               "W" writeable, "A" aligned, "E" little-endian

Functional-style: apply_along_axis and vectorize

# Apply function along one axis
def f(row):
    return row.max() - row.min()

b = np.arange(12).reshape(3, 4)
np.apply_along_axis(f, axis=1, arr=b)   # apply f to each row → (3,)
np.apply_along_axis(f, axis=0, arr=b)   # apply f to each col → (4,)

# np.apply_over_axes — reduce over multiple axes
np.apply_over_axes(np.sum, b, axes=[0, 1])   # sum over both axes

# Vectorize a scalar function
def safe_div(x, y):
    return x / y if y != 0 else 0.0

vf = np.vectorize(safe_div)
vf(np.arange(5), np.array([1, 0, 2, 0, 4]))   # [0., 0., 1., 0., 1.]

# excluded — arguments that should NOT be vectorized
vf = np.vectorize(safe_div, excluded=["y"])

np.vectorize is a convenience wrapper around Python loops — it does not give C-speed. For performance, use ufuncs or np.where.

Padding

a = np.array([1, 2, 3])
np.pad(a, pad_width=2, mode="constant", constant_values=0)
# [0, 0, 1, 2, 3, 0, 0]

np.pad(a, (2, 3), mode="constant", constant_values=(10, 20))
# [10, 10, 1, 2, 3, 20, 20, 20]

b = np.array([[1, 2], [3, 4]])
np.pad(b, 1, mode="constant")    # 1-wide zero border
np.pad(b, 1, mode="reflect")     # mirror padding
np.pad(b, 1, mode="edge")        # replicate edge values
np.pad(b, 1, mode="wrap")        # periodic padding
np.pad(b, 1, mode="symmetric")   # reflect about edge (includes edge)
np.pad(b, ((1, 2), (3, 4)))      # per-axis (before, after) widths

Interpolation

xp = np.array([0.0, 1.0, 2.0, 3.0])
fp = np.array([0.0, 1.0, 4.0, 9.0])

np.interp(1.5, xp, fp)                     # 2.5 — linear interpolation
np.interp([0.5, 1.5, 2.5], xp, fp)        # [0.5, 2.5, 6.5]
np.interp(x, xp, fp, left=None, right=None)  # values outside xp range
np.interp(-1.0, xp, fp, left=-999)          # extrapolation value

Piecewise Functions

x = np.linspace(-2, 2, 9)
np.piecewise(
    x,
    [x < 0, x == 0, x > 0],
    [-1, 0, 1]                # constants or callables per condition
)

np.piecewise(
    x,
    [x < 0, x >= 0],
    [lambda v: -v, lambda v: v]   # abs(x)
)

Convolution and Correlation

a = np.array([1, 2, 3, 4, 5])
b = np.array([1, 0, -1])

np.convolve(a, b)                  # mode="full" default → length len(a)+len(b)-1
np.convolve(a, b, mode="same")    # same length as a
np.convolve(a, b, mode="valid")   # only where they fully overlap

np.correlate(a, b)                 # cross-correlation
np.correlate(a, b, mode="full")

Windowing Functions (for signal processing)

np.bartlett(10)        # Bartlett window
np.blackman(10)        # Blackman window
np.hamming(10)         # Hamming window
np.hanning(10)         # Hanning window
np.kaiser(10, beta=14) # Kaiser window

# Usage: multiply signal by window before FFT
signal = np.random.randn(64)
window = np.hanning(64)
windowed = signal * window

Utilities for Arrays

# Stack / unpack multiple outputs
np.column_stack([[1, 2], [3, 4]])  # 2×2 array (1-D arrays become columns)
np.vstack([[1, 2], [3, 4]])        # stack as rows (np.row_stack is deprecated in NumPy 2.0+)

np.atleast_1d(3)           # scalar → array([3])
np.atleast_1d([1, 2])      # list   → array([1, 2])
np.atleast_2d([1, 2])      # (1, 2) → column vector
np.atleast_3d([1, 2])      # → (1, 1, 2)

# Check shapes / consistency
np.broadcast_shapes((3,), (3, 1), (1, 3))   # (3, 3)
np.shape(5)                # () — shape of scalar

# Array iteration helpers
for row in b:              # iterates over rows of 2-D
    pass
for x in np.nditer(b):    # element-by-element, C order
    pass
for x in b.flat:          # flat iterator
    pass

# nditer for multi-array operations
it = np.nditer([a, b], flags=["buffered"], op_flags=[["readonly"], ["readwrite"]])
with it:
    for x, y in it:
        y[...] = x * 2

Input / Output Utilities

np.array_str(a, precision=4)     # string repr
np.array_repr(a)                  # repr string
np.array2string(a, separator=", ", prefix="arr=")

# Binary file I/O
np.save("file.npy", a)
a = np.load("file.npy")
np.savez("archive.npz", a=a, b=b)
npz = np.load("archive.npz")
npz.files        # ["a", "b"]
npz["a"]

# Text I/O
np.savetxt("file.csv", a, delimiter=",", fmt="%.4f", header="x,y,z")
np.loadtxt("file.csv", delimiter=",", skiprows=1)

Miscellaneous

np.bartlett(10)           # signal windows (see above)
np.sort(a, axis=0)        # (np.msort was removed in NumPy 2.0)
np.array_split(a, 3)      # split into (possibly unequal) parts
np.dsplit(c, 2)           # depth-split 3-D arrays

np.ix_([0, 1], [2, 3])   # open mesh from index arrays → for outer indexing

np.ndindex(2, 3)          # iterator over (i, j) multi-indices
list(np.ndindex(2, 3))    # [(0,0),(0,1),(0,2),(1,0),(1,1),(1,2)]

np.ndenumerate(b)         # iterator of (index_tuple, value)
for idx, val in np.ndenumerate(b):
    pass

np.unravel_index([5, 11], (3, 4))   # ([1, 2], [1, 3])  flat→multi index
np.ravel_multi_index(([1, 2], [1, 3]), (3, 4))  # [5, 11]

np.shares_memory(a, b)   # True if a and b overlap in memory
np.may_share_memory(a, b)  # faster but may have false positives