NumPy Cheatsheet

Useful Functions

Use this NumPy reference while you build software engineering projects, review code for technical interview prep, or polish examples for a software engineer resume.

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