PythonMastery
intermediate 22 min read · lesson 1 of 9 in Data Science & ML

NumPy: The Numerical Engine

1 · The lesson

read

Python's built-in lists are arrays of pointers to objects of any type. They are flexible, but very slow for numerical work. NumPy uses a different design. An ndarray is a single block of contiguous memory, where every element has the same fixed-width type, plus a small header that describes the shape and strides. Because the data is stored this way, NumPy can use vectorised C loops, SIMD instructions and linear algebra through BLAS. The result is often 50–200× faster than the same code in plain Python.

You've seen the basics in ds-numpy. This lesson goes further: memory layout, broadcasting rules, advanced indexing, how axis works, and the performance habits that make NumPy code fast enough for production. By the end, you should write array expressions by default, instead of a for loop over rows.


1. The Memory Model: Why NumPy Is Fast

An ndarray is three things bundled together:

1. A data buffer: a flat C array of values that all have the same width.
2. A dtype: the element type (int32, float64, bool, etc.) and its byte width.
3. A shape + strides tuple: how to read the buffer as an n-dimensional grid.

python
import numpy as np

a = np.arange(12, dtype=np.int32).reshape(3, 4)
print(a)
# [[ 0  1  2  3]
#  [ 4  5  6  7]
#  [ 8  9 10 11]]

print(a.dtype)         # int32
print(a.shape)         # (3, 4)
print(a.strides)       # (16, 4)   16 bytes to the next row, 4 bytes to the next column
print(a.nbytes)        # 48        12 ints * 4 bytes each

strides=(16, 4) is the important part. Moving down one row jumps 16 bytes (one full row of four int32s). Moving across one column jumps 4 bytes (one int32). Reshaping, slicing and transposing usually only change the strides, and no data is copied. That's why a.T is instantaneous on a 1 GB array.

Compare this with a Python list of lists. Each row is a separate PyListObject, each cell is a separate PyLongObject somewhere else in memory, and every access follows several pointers. NumPy keeps the data close together in memory, and CPUs process data like that much faster.


2. Creating Arrays Efficiently

Avoid np.array([0]*1000), because it builds a Python list first. The right constructors allocate the buffer directly:

python
import numpy as np

np.zeros((3, 4), dtype=np.float32)      # filled with 0.0, 12 floats = 48 bytes
np.ones((2, 2), dtype=np.int64)         # filled with 1
np.full((2, 3), 7.5)                    # filled with 7.5
np.empty((3, 3))                        # memory is not cleared, so overwrite every value before use
np.eye(4)                               # 4x4 identity matrix
np.arange(0, 10, 0.5)                   # like range, but allows a float step and returns an array
np.linspace(0, 1, 11)                   # 11 evenly spaced points from 0 to 1 inclusive

dtype choice matters for memory. A (1000, 1000) float64 array is 8 MB. The same shape as float32 is 4 MB, and as int8 it's 1 MB. Pick the smallest type that can hold your data, but only when you have a reason to. Integer overflow happens without any warning (see Common Mistakes).

For random data, use the newer Generator API, not the legacy np.random.rand / randn functions with their global state:

python
rng = np.random.default_rng(seed=42)     # the modern API (NEP 19)
rng.standard_normal((3, 4))              # 3x4 array from the standard normal N(0, 1)
rng.uniform(0, 100, size=1000)           # 1000 uniform floats in [0, 100)
rng.integers(0, 10, size=20)             # 20 ints in [0, 10)
rng.choice(["A", "B", "C"], size=5, p=[0.5, 0.3, 0.2])
+ setup added so this can run · re-runs 2 earlier blocks on this page first
# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

np = _AutoMock('np')

# Earlier blocks on this page that this one builds on, re-run quietly first.
try:
    import numpy as np
except Exception:
    pass
_pm_earlier = [
    'import numpy as np\n\na = np.arange(12, dtype=np.int32).reshape(3, 4)\nprint(a)\n# [[ 0  1  2  3]\n#  [ 4  5  6  7]\n#  [ 8  9 10 11]]\n\nprint(a.dtype)         # int32\nprint(a.shape)         # (3, 4)\nprint(a.strides)       # (16, 4)   16 bytes to the next row, 4 bytes to the next column\nprint(a.nbytes)        # 48        12 ints * 4 bytes each',
    'import numpy as np\n\nnp.zeros((3, 4), dtype=np.float32)      # filled with 0.0, 12 floats = 48 bytes\nnp.ones((2, 2), dtype=np.int64)         # filled with 1\nnp.full((2, 3), 7.5)                    # filled with 7.5\nnp.empty((3, 3))                        # memory is not cleared, so overwrite every value before use\nnp.eye(4)                               # 4x4 identity matrix\nnp.arange(0, 10, 0.5)                   # like range, but allows a float step and returns an array\nnp.linspace(0, 1, 11)                   # 11 evenly spaced points from 0 to 1 inclusive',
]
def _pm_run_earlier(_srcs, _g):
    import sys as _s, io as _io, warnings as _w
    for _src in _srcs:
        _o, _e = _s.stdout, _s.stderr
        _s.stdout = _s.stderr = _io.StringIO()
        try:
            with _w.catch_warnings():
                _w.simplefilter('ignore')
                exec(compile(_src, '<earlier block>', 'exec'), _g)
        except (Exception, SystemExit):
            pass
        finally:
            _s.stdout, _s.stderr = _o, _e
    if 'matplotlib.pyplot' in _s.modules:
        _s.modules['matplotlib.pyplot'].close('all')   # their plots belong to their own blocks
_pm_run_earlier(_pm_earlier, globals())
del _pm_earlier, _pm_run_earlier

The legacy functions still work, but they share one global seed, which causes problems when code runs in parallel. Use default_rng in new code.


3. Vectorisation: Whole Arrays, Not Loops

The main rule for writing fast NumPy: write your computation as an operation on whole arrays, never as a Python loop over each element. Element-wise arithmetic, comparisons and unary functions are all vectorised:

python
import numpy as np

a = np.array([1, 2, 3, 4])
b = np.array([10, 20, 30, 40])

a + b                       # [11, 22, 33, 44]   element by element
a * b                       # [10, 40, 90, 160]
a ** 2                      # [1, 4, 9, 16]
np.sqrt(a)                  # [1., 1.41, 1.73, 2.]
np.exp(a)                   # [2.72, 7.39, 20.09, 54.60]
a > 2                       # [False, False, True, True]   a boolean array

# The "I want a Python loop" trap
result = np.empty_like(a)
for i in range(len(a)):
    result[i] = a[i] ** 2 + b[i]      # SLOW: a Python loop over a C array

# The vectorised way
result = a ** 2 + b                   # FAST: one pass in compiled C code

A vectorised expression on a million-element array runs in milliseconds. The same work in a Python loop can take seconds. Every operator and every function in np.* is built to work on whole arrays at once.


4. Broadcasting in Depth

Broadcasting is the rule that lets arrays of different shapes work together without you copying data to make the shapes match. NumPy lines up the shapes from the right. Then, for each axis, one of three things must be true: the sizes are equal, one of them is 1 (and is stretched), or one is missing (and is treated as 1).

The rules

OperationResult shapeReasoning
(3, 4) + (4,)(3, 4)RHS becomes (1, 4), stretched to (3, 4)
(3, 1) * (1, 4)(3, 4)Both stretched to (3, 4): outer product
(3, 4) + (3,)errorRight-align: 4 != 3 and neither is 1
(3, 4) + (3, 1)(3, 4)Column vector stretched across columns
(5,) + 5.0(5,)Scalar broadcasts to anything

No memory is used for the "stretching". NumPy reads the broadcast axis with a stride of zero, so the same values are reused without being copied.

Common patterns

python
import numpy as np

# Subtract the column mean from each column (centre per feature)
X = np.array([[1.0, 2.0, 3.0],
              [4.0, 5.0, 6.0],
              [7.0, 8.0, 9.0]])

col_mean = X.mean(axis=0)               # shape (3,)
centred  = X - col_mean                 # (3,3) - (3,) gives (3,3), and each column now has mean 0

# Z-score normalisation per column
std = X.std(axis=0)
z = (X - col_mean) / std

# Outer product via broadcasting (no np.outer needed)
u = np.array([1, 2, 3])                 # shape (3,)
v = np.array([10, 20, 30, 40])          # shape (4,)
outer = u[:, None] * v[None, :]         # (3,1) * (1,4) gives (3,4)
# [[ 10  20  30  40]
#  [ 20  40  60  80]
#  [ 30  60  90 120]]

# Pairwise distances (Euclidean) between rows of A and rows of B
A = np.random.default_rng(0).standard_normal((100, 5))
B = np.random.default_rng(1).standard_normal((50, 5))
# (100,1,5) - (1,50,5) gives (100,50,5). Summing squares on the last axis gives (100,50)
dists = np.sqrt(((A[:, None, :] - B[None, :, :]) ** 2).sum(axis=-1))
print(dists.shape)                      # (100, 50)

X[:, None] and X[None, :] are the idioms you will use most. They add a new axis of length 1, so that broadcasting can do the rest.


5. Advanced Indexing

Besides a[0] and a[1:5], there are three more flexible ways to index.

Integer-array indexing: pick any rows or columns

python
import numpy as np

a = np.arange(20).reshape(5, 4)
# [[ 0  1  2  3]
#  [ 4  5  6  7]
#  [ 8  9 10 11]
#  [12 13 14 15]
#  [16 17 18 19]]

a[[0, 2, 4]]                # rows 0, 2 and 4, shape (3, 4)
a[:, [3, 1, 0]]             # columns in a new order, shape (5, 3)
a[[0, 1, 2], [3, 2, 1]]     # picks (0,3), (1,2) and (2,1), giving [3, 6, 9]

Boolean masks: filter by condition

python
mask = a > 10
a[mask]                     # 1-D array of every element > 10: [11, 12, 13, 14, 15, 16, 17, 18, 19]
a[a % 2 == 0]               # all even values

# Mask assignment: set every negative entry to zero
x = np.array([-3, 1, -1, 5, -2])
x[x < 0] = 0
print(x)                    # [0 1 0 5 0]
+ setup added so this can run · re-runs 6 earlier blocks on this page first
# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

a = _AutoMock('a')
np = _AutoMock('np')

# Earlier blocks on this page that this one builds on, re-run quietly first.
try:
    import numpy as np
except Exception:
    pass
_pm_earlier = [
    'import numpy as np\n\na = np.arange(12, dtype=np.int32).reshape(3, 4)\nprint(a)\n# [[ 0  1  2  3]\n#  [ 4  5  6  7]\n#  [ 8  9 10 11]]\n\nprint(a.dtype)         # int32\nprint(a.shape)         # (3, 4)\nprint(a.strides)       # (16, 4)   16 bytes to the next row, 4 bytes to the next column\nprint(a.nbytes)        # 48        12 ints * 4 bytes each',
    'import numpy as np\n\nnp.zeros((3, 4), dtype=np.float32)      # filled with 0.0, 12 floats = 48 bytes\nnp.ones((2, 2), dtype=np.int64)         # filled with 1\nnp.full((2, 3), 7.5)                    # filled with 7.5\nnp.empty((3, 3))                        # memory is not cleared, so overwrite every value before use\nnp.eye(4)                               # 4x4 identity matrix\nnp.arange(0, 10, 0.5)                   # like range, but allows a float step and returns an array\nnp.linspace(0, 1, 11)                   # 11 evenly spaced points from 0 to 1 inclusive',
    'rng = np.random.default_rng(seed=42)     # the modern API (NEP 19)\nrng.standard_normal((3, 4))              # 3x4 array from the standard normal N(0, 1)\nrng.uniform(0, 100, size=1000)           # 1000 uniform floats in [0, 100)\nrng.integers(0, 10, size=20)             # 20 ints in [0, 10)\nrng.choice(["A", "B", "C"], size=5, p=[0.5, 0.3, 0.2])',
    'import numpy as np\n\na = np.array([1, 2, 3, 4])\nb = np.array([10, 20, 30, 40])\n\na + b                       # [11, 22, 33, 44]   element by element\na * b                       # [10, 40, 90, 160]\na ** 2                      # [1, 4, 9, 16]\nnp.sqrt(a)                  # [1., 1.41, 1.73, 2.]\nnp.exp(a)                   # [2.72, 7.39, 20.09, 54.60]\na > 2                       # [False, False, True, True]   a boolean array\n\n# The "I want a Python loop" trap\nresult = np.empty_like(a)\nfor i in range(len(a)):\n    result[i] = a[i] ** 2 + b[i]      # SLOW: a Python loop over a C array\n\n# The vectorised way\nresult = a ** 2 + b                   # FAST: one pass in compiled C code',
    'import numpy as np\n\n# Subtract the column mean from each column (centre per feature)\nX = np.array([[1.0, 2.0, 3.0],\n              [4.0, 5.0, 6.0],\n              [7.0, 8.0, 9.0]])\n\ncol_mean = X.mean(axis=0)               # shape (3,)\ncentred  = X - col_mean                 # (3,3) - (3,) gives (3,3), and each column now has mean 0\n\n# Z-score normalisation per column\nstd = X.std(axis=0)\nz = (X - col_mean) / std\n\n# Outer product via broadcasting (no np.outer needed)\nu = np.array([1, 2, 3])                 # shape (3,)\nv = np.array([10, 20, 30, 40])          # shape (4,)\nouter = u[:, None] * v[None, :]         # (3,1) * (1,4) gives (3,4)\n# [[ 10  20  30  40]\n#  [ 20  40  60  80]\n#  [ 30  60  90 120]]\n\n# Pairwise distances (Euclidean) between rows of A and rows of B\nA = np.random.default_rng(0).standard_normal((100, 5))\nB = np.random.default_rng(1).standard_normal((50, 5))\n# (100,1,5) - (1,50,5) gives (100,50,5). Summing squares on the last axis gives (100,50)\ndists = np.sqrt(((A[:, None, :] - B[None, :, :]) ** 2).sum(axis=-1))\nprint(dists.shape)                      # (100, 50)',
    'import numpy as np\n\na = np.arange(20).reshape(5, 4)\n# [[ 0  1  2  3]\n#  [ 4  5  6  7]\n#  [ 8  9 10 11]\n#  [12 13 14 15]\n#  [16 17 18 19]]\n\na[[0, 2, 4]]                # rows 0, 2 and 4, shape (3, 4)\na[:, [3, 1, 0]]             # columns in a new order, shape (5, 3)\na[[0, 1, 2], [3, 2, 1]]     # picks (0,3), (1,2) and (2,1), giving [3, 6, 9]',
]
def _pm_run_earlier(_srcs, _g):
    import sys as _s, io as _io, warnings as _w
    for _src in _srcs:
        _o, _e = _s.stdout, _s.stderr
        _s.stdout = _s.stderr = _io.StringIO()
        try:
            with _w.catch_warnings():
                _w.simplefilter('ignore')
                exec(compile(_src, '<earlier block>', 'exec'), _g)
        except (Exception, SystemExit):
            pass
        finally:
            _s.stdout, _s.stderr = _o, _e
    if 'matplotlib.pyplot' in _s.modules:
        _s.modules['matplotlib.pyplot'].close('all')   # their plots belong to their own blocks
_pm_run_earlier(_pm_earlier, globals())
del _pm_earlier, _pm_run_earlier

np.where: a vectorised if/else

python
x = np.array([-3, 1, -1, 5, -2])
np.where(x < 0, 0, x)       # [0, 1, 0, 5, 0]   the same result as np.maximum(x, 0)
np.where(x < 0, "neg", "pos")
# array(['neg', 'pos', 'neg', 'pos', 'neg'], dtype='<U3')

# The one-argument form returns the indices where the condition is True
idx = np.where(x < 0)
print(idx)                  # (array([0, 2, 4]),)
print(x[idx])               # [-3 -1 -2]
+ setup added so this can run · re-runs 7 earlier blocks on this page first
# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

np = _AutoMock('np')
a = _AutoMock('a')

# Earlier blocks on this page that this one builds on, re-run quietly first.
try:
    import numpy as np
except Exception:
    pass
_pm_earlier = [
    'import numpy as np\n\na = np.arange(12, dtype=np.int32).reshape(3, 4)\nprint(a)\n# [[ 0  1  2  3]\n#  [ 4  5  6  7]\n#  [ 8  9 10 11]]\n\nprint(a.dtype)         # int32\nprint(a.shape)         # (3, 4)\nprint(a.strides)       # (16, 4)   16 bytes to the next row, 4 bytes to the next column\nprint(a.nbytes)        # 48        12 ints * 4 bytes each',
    'import numpy as np\n\nnp.zeros((3, 4), dtype=np.float32)      # filled with 0.0, 12 floats = 48 bytes\nnp.ones((2, 2), dtype=np.int64)         # filled with 1\nnp.full((2, 3), 7.5)                    # filled with 7.5\nnp.empty((3, 3))                        # memory is not cleared, so overwrite every value before use\nnp.eye(4)                               # 4x4 identity matrix\nnp.arange(0, 10, 0.5)                   # like range, but allows a float step and returns an array\nnp.linspace(0, 1, 11)                   # 11 evenly spaced points from 0 to 1 inclusive',
    'rng = np.random.default_rng(seed=42)     # the modern API (NEP 19)\nrng.standard_normal((3, 4))              # 3x4 array from the standard normal N(0, 1)\nrng.uniform(0, 100, size=1000)           # 1000 uniform floats in [0, 100)\nrng.integers(0, 10, size=20)             # 20 ints in [0, 10)\nrng.choice(["A", "B", "C"], size=5, p=[0.5, 0.3, 0.2])',
    'import numpy as np\n\na = np.array([1, 2, 3, 4])\nb = np.array([10, 20, 30, 40])\n\na + b                       # [11, 22, 33, 44]   element by element\na * b                       # [10, 40, 90, 160]\na ** 2                      # [1, 4, 9, 16]\nnp.sqrt(a)                  # [1., 1.41, 1.73, 2.]\nnp.exp(a)                   # [2.72, 7.39, 20.09, 54.60]\na > 2                       # [False, False, True, True]   a boolean array\n\n# The "I want a Python loop" trap\nresult = np.empty_like(a)\nfor i in range(len(a)):\n    result[i] = a[i] ** 2 + b[i]      # SLOW: a Python loop over a C array\n\n# The vectorised way\nresult = a ** 2 + b                   # FAST: one pass in compiled C code',
    'import numpy as np\n\n# Subtract the column mean from each column (centre per feature)\nX = np.array([[1.0, 2.0, 3.0],\n              [4.0, 5.0, 6.0],\n              [7.0, 8.0, 9.0]])\n\ncol_mean = X.mean(axis=0)               # shape (3,)\ncentred  = X - col_mean                 # (3,3) - (3,) gives (3,3), and each column now has mean 0\n\n# Z-score normalisation per column\nstd = X.std(axis=0)\nz = (X - col_mean) / std\n\n# Outer product via broadcasting (no np.outer needed)\nu = np.array([1, 2, 3])                 # shape (3,)\nv = np.array([10, 20, 30, 40])          # shape (4,)\nouter = u[:, None] * v[None, :]         # (3,1) * (1,4) gives (3,4)\n# [[ 10  20  30  40]\n#  [ 20  40  60  80]\n#  [ 30  60  90 120]]\n\n# Pairwise distances (Euclidean) between rows of A and rows of B\nA = np.random.default_rng(0).standard_normal((100, 5))\nB = np.random.default_rng(1).standard_normal((50, 5))\n# (100,1,5) - (1,50,5) gives (100,50,5). Summing squares on the last axis gives (100,50)\ndists = np.sqrt(((A[:, None, :] - B[None, :, :]) ** 2).sum(axis=-1))\nprint(dists.shape)                      # (100, 50)',
    'import numpy as np\n\na = np.arange(20).reshape(5, 4)\n# [[ 0  1  2  3]\n#  [ 4  5  6  7]\n#  [ 8  9 10 11]\n#  [12 13 14 15]\n#  [16 17 18 19]]\n\na[[0, 2, 4]]                # rows 0, 2 and 4, shape (3, 4)\na[:, [3, 1, 0]]             # columns in a new order, shape (5, 3)\na[[0, 1, 2], [3, 2, 1]]     # picks (0,3), (1,2) and (2,1), giving [3, 6, 9]',
    'mask = a > 10\na[mask]                     # 1-D array of every element > 10: [11, 12, 13, 14, 15, 16, 17, 18, 19]\na[a % 2 == 0]               # all even values\n\n# Mask assignment: set every negative entry to zero\nx = np.array([-3, 1, -1, 5, -2])\nx[x < 0] = 0\nprint(x)                    # [0 1 0 5 0]',
]
def _pm_run_earlier(_srcs, _g):
    import sys as _s, io as _io, warnings as _w
    for _src in _srcs:
        _o, _e = _s.stdout, _s.stderr
        _s.stdout = _s.stderr = _io.StringIO()
        try:
            with _w.catch_warnings():
                _w.simplefilter('ignore')
                exec(compile(_src, '<earlier block>', 'exec'), _g)
        except (Exception, SystemExit):
            pass
        finally:
            _s.stdout, _s.stderr = _o, _e
    if 'matplotlib.pyplot' in _s.modules:
        _s.modules['matplotlib.pyplot'].close('all')   # their plots belong to their own blocks
_pm_run_earlier(_pm_earlier, globals())
del _pm_earlier, _pm_run_earlier

Fancy indexing always creates a copy, never a view (more on that below). Boolean masks also copy. Plain slices (a[1:5]) are views, so changing them changes the original.


6. Aggregations and the axis Argument

Reductions like sum, mean, std, min, max, argmax all accept axis=. The idea: axis=k removes axis k and keeps the others.

python
import numpy as np

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

a.sum()                       # 21    adds up every element
a.sum(axis=0)                 # [5, 7, 9]   collapses the rows, giving one sum per column
a.sum(axis=1)                 # [6, 15]     collapses the columns, giving one sum per row

a.mean(axis=0)                # [2.5, 3.5, 4.5]   the column means
a.argmax(axis=1)              # [2, 2]            the index of the max in each row

A way to remember it: axis=0 goes "down through the rows" and gives one result per column. axis=1 goes "across the columns" and gives one result per row. It feels backwards at first, but you get used to it.

keepdims=True keeps the reduced axis with length 1. You need this to broadcast the result back onto the original array:

python
row_max = a.max(axis=1)                  # shape (2,)   cannot be subtracted from (2,3)
row_max = a.max(axis=1, keepdims=True)   # shape (2, 1) (this one broadcasts correctly)
normalised = a - row_max                 # subtract each row's max (the trick for a stable softmax)
+ setup added so this can run · re-runs 4 earlier blocks on this page first
# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

a = _AutoMock('a')

# Earlier blocks on this page that this one builds on, re-run quietly first.
try:
    import numpy as np
except Exception:
    pass
_pm_earlier = [
    'import numpy as np\n\na = np.arange(12, dtype=np.int32).reshape(3, 4)\nprint(a)\n# [[ 0  1  2  3]\n#  [ 4  5  6  7]\n#  [ 8  9 10 11]]\n\nprint(a.dtype)         # int32\nprint(a.shape)         # (3, 4)\nprint(a.strides)       # (16, 4)   16 bytes to the next row, 4 bytes to the next column\nprint(a.nbytes)        # 48        12 ints * 4 bytes each',
    'import numpy as np\n\na = np.array([1, 2, 3, 4])\nb = np.array([10, 20, 30, 40])\n\na + b                       # [11, 22, 33, 44]   element by element\na * b                       # [10, 40, 90, 160]\na ** 2                      # [1, 4, 9, 16]\nnp.sqrt(a)                  # [1., 1.41, 1.73, 2.]\nnp.exp(a)                   # [2.72, 7.39, 20.09, 54.60]\na > 2                       # [False, False, True, True]   a boolean array\n\n# The "I want a Python loop" trap\nresult = np.empty_like(a)\nfor i in range(len(a)):\n    result[i] = a[i] ** 2 + b[i]      # SLOW: a Python loop over a C array\n\n# The vectorised way\nresult = a ** 2 + b                   # FAST: one pass in compiled C code',
    'import numpy as np\n\na = np.arange(20).reshape(5, 4)\n# [[ 0  1  2  3]\n#  [ 4  5  6  7]\n#  [ 8  9 10 11]\n#  [12 13 14 15]\n#  [16 17 18 19]]\n\na[[0, 2, 4]]                # rows 0, 2 and 4, shape (3, 4)\na[:, [3, 1, 0]]             # columns in a new order, shape (5, 3)\na[[0, 1, 2], [3, 2, 1]]     # picks (0,3), (1,2) and (2,1), giving [3, 6, 9]',
    'import numpy as np\n\na = np.array([[1, 2, 3],\n              [4, 5, 6]])     # shape (2, 3)\n\na.sum()                       # 21    adds up every element\na.sum(axis=0)                 # [5, 7, 9]   collapses the rows, giving one sum per column\na.sum(axis=1)                 # [6, 15]     collapses the columns, giving one sum per row\n\na.mean(axis=0)                # [2.5, 3.5, 4.5]   the column means\na.argmax(axis=1)              # [2, 2]            the index of the max in each row',
]
def _pm_run_earlier(_srcs, _g):
    import sys as _s, io as _io, warnings as _w
    for _src in _srcs:
        _o, _e = _s.stdout, _s.stderr
        _s.stdout = _s.stderr = _io.StringIO()
        try:
            with _w.catch_warnings():
                _w.simplefilter('ignore')
                exec(compile(_src, '<earlier block>', 'exec'), _g)
        except (Exception, SystemExit):
            pass
        finally:
            _s.stdout, _s.stderr = _o, _e
    if 'matplotlib.pyplot' in _s.modules:
        _s.modules['matplotlib.pyplot'].close('all')   # their plots belong to their own blocks
_pm_run_earlier(_pm_earlier, globals())
del _pm_earlier, _pm_run_earlier

7. Linear Algebra: @, dot, solve

@ is the matrix-multiplication operator (PEP 465). For arrays with more than two dimensions, it multiplies a batch of matrices at once, using broadcasting.

python
import numpy as np

A = np.array([[1, 2], [3, 4]], dtype=float)
b = np.array([5, 6], dtype=float)

A @ b                           # [17., 39.]    matrix times vector
A @ A                           # matrix times matrix, a 2x2 result
A.T @ A                         # Gram matrix

# Solve Ax = b. Do NOT compute inv(A) @ b
x = np.linalg.solve(A, b)       # [-4., 4.5]    faster and more accurate than the inverse

# When you really need the inverse
A_inv = np.linalg.inv(A)

# Norms: Euclidean, Manhattan and Frobenius
np.linalg.norm(b)               # 7.81   L2 (Euclidean) by default
np.linalg.norm(b, ord=1)        # 11.0   L1 (Manhattan)
np.linalg.norm(A, ord="fro")    # 5.48   Frobenius

# Eigendecomposition
eigvals, eigvecs = np.linalg.eig(A)
# eigvals: [-0.37,  5.37]
# eigvecs: columns are the corresponding unit eigenvectors

# Singular value decomposition (the method behind PCA)
U, s, Vt = np.linalg.svd(A, full_matrices=False)

For Ax = b, always use solve instead of inv @ b. It uses LU decomposition internally, it is faster, and it avoids the rounding problems of computing the inverse.


8. Statistics and Random: Beyond the Basics

python
import numpy as np

data = np.array([2, 4, 4, 4, 5, 5, 7, 9])

np.percentile(data, [25, 50, 75])       # [4., 4.5, 5.5]   the quartiles
np.median(data)                         # 4.5
np.cov(data)                            # 4.57   sample covariance (for 1-D data, the variance)

# Correlation matrix
x = np.random.default_rng(0).standard_normal((100,))
y = 2 * x + 0.5 * np.random.default_rng(1).standard_normal((100,))
np.corrcoef(x, y)
# [[1.   0.97]
#  [0.97 1.  ]]

# Statistics on 2-D data: choose the axis carefully
X = np.random.default_rng(0).standard_normal((1000, 5))   # 1000 samples, 5 features
feature_mean = X.mean(axis=0)           # shape (5,)   one mean per feature
feature_std  = X.std(axis=0, ddof=1)    # sample std (N-1 denominator)

ddof=1 gives the unbiased sample statistic. NumPy defaults to the population version (ddof=0), but pandas defaults to the sample version. This difference causes many mistakes.


9. Reshape, Transpose, Stack, Split

Reshaping, transposing and splitting usually don't copy data. They only read the same buffer in a different way. Stacking is different: it always creates a new array.

python
import numpy as np

a = np.arange(12)
a.reshape(3, 4)                 # 3x4 view
a.reshape(2, -1)                # -1 means "work it out", so 2x6
a.reshape(2, 2, 3)              # 3-D view

# Transpose: swap or reorder the axes
M = np.arange(24).reshape(2, 3, 4)
M.T.shape                       # (4, 3, 2)   reverses all axes
M.transpose(1, 0, 2).shape      # (3, 2, 4)   axes in the order you give

# Stacking and concatenation
a = np.array([1, 2, 3])
b = np.array([4, 5, 6])

np.concatenate([a, b])          # [1, 2, 3, 4, 5, 6]
np.stack([a, b])                # [[1,2,3],[4,5,6]]   adds a new axis at the front
np.vstack([a, b])               # [[1,2,3],[4,5,6]]   stacks them as rows
np.hstack([a, b])               # [1, 2, 3, 4, 5, 6]  joins 1-D arrays end to end
np.column_stack([a, b])         # [[1,4],[2,5],[3,6]] (each array becomes a column)

# Splitting: the opposite of stacking
M = np.arange(12).reshape(3, 4)
np.split(M, 2, axis=1)          # split into 2 along the columns, giving a list of (3,2) arrays

reshape requires the new shape to have the same total size. ravel() flattens to 1-D (a view if the data is contiguous). flatten() always copies.


10. Views vs Copies: The Trap

Slicing returns a view: the same buffer with a different shape or stride. Changing a view changes the original. Fancy indexing returns a copy: separate data.

python
import numpy as np

a = np.arange(10)
b = a[2:6]                      # VIEW: shares memory with a
b[0] = 999
print(a)                        # [  0   1 999   3   4   5   6   7   8   9]   a changed too!

c = a[[2, 3, 4]]                # COPY: fancy indexing always copies
c[0] = -1
print(a)                        # unchanged

# Check explicitly
print(b.base is a)              # True   b is a view of a
print(c.base is a)              # False  c is a separate copy

# Force a copy when you want one
d = a[2:6].copy()
d[0] = 555
print(a)                        # unchanged

This causes many "why did my data change?" bugs. When you plan to change the data, call .copy() first to be safe. When you want a view (for memory or speed), check it with .base.


11. Performance: Benchmarks, np.einsum and the Cost of Loops

This benchmark answers the question "should I vectorise?":

python
import numpy as np, time

n = 1_000_000
a = np.random.default_rng(0).standard_normal(n)
b = np.random.default_rng(1).standard_normal(n)

# Python loop
t0 = time.perf_counter()
result_loop = [a[i] * b[i] + a[i] for i in range(n)]
t_loop = time.perf_counter() - t0

# NumPy vectorised
t0 = time.perf_counter()
result_vec = a * b + a
t_vec = time.perf_counter() - t0

print(f"loop: {t_loop:.3f}s   vec: {t_vec:.4f}s   speedup: {t_loop/t_vec:.0f}x")
# loop: ~0.300s   vec: ~0.003s   speedup: ~100x

For complex multi-axis operations, np.einsum expresses tensor contractions in Einstein notation. It is sometimes faster than equivalent combinations of transpose, reshape, and @:

python
import numpy as np

A = np.random.default_rng(0).standard_normal((100, 50))
B = np.random.default_rng(1).standard_normal((50, 30))

# Standard matmul
np.einsum("ij,jk->ik", A, B).shape          # (100, 30)   the same as A @ B

# Batched dot product of corresponding rows
X = np.random.default_rng(0).standard_normal((1000, 5))
Y = np.random.default_rng(1).standard_normal((1000, 5))
dots = np.einsum("ij,ij->i", X, Y)          # shape (1000,) (one dot product per row)

When NumPy's vectorised operations aren't enough, for example a busy inner loop with branches that can't be written as array operations, use numba (@njit JIT-compiles a Python function) or Cython. Both run that loop as compiled code instead of in the Python interpreter.


12. I/O: Saving and Loading

python
import numpy as np

a = np.arange(12).reshape(3, 4)

np.save("a.npy", a)                         # binary, fastest, and keeps the dtype
b = np.load("a.npy")

np.savez("bundle.npz", x=a, y=a * 2)        # multiple arrays in one file
data = np.load("bundle.npz")
print(data["x"], data["y"])

np.savetxt("a.csv", a, delimiter=",", fmt="%d")   # human-readable, but can lose float precision
b = np.loadtxt("a.csv", delimiter=",")

.npy is the default for "save this for later". It loads back exactly the same array, including dtype, shape and byte order. .csv is for sharing data with people or tools that don't use NumPy.


Common Mistakes

1. Silent integer overflow

python
import numpy as np

a = np.array([100, 100], dtype=np.int8)     # range is [-128, 127], so 100 fits
print(a + a)                                # [-56 -56]  it WRAPPED around with no warning

NumPy integer dtypes wrap around on overflow without any warning. int8, int16 and int32 can all do this. Use int64 (the default on 64-bit platforms) unless you are short of memory. If you use a smaller type, check the range of the values when the data comes in.

2. Mixed dtypes producing object arrays

python
a = np.array([1, 2, "three", 4])
print(a.dtype)                              # <U21   Unicode strings, which are slow

a = np.array([1, 2, None, 4])
print(a.dtype)                              # object  separate Python objects, the slowest
+ setup added so this can run · re-runs 17 earlier blocks on this page first
# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

np = _AutoMock('np')
a = _AutoMock('a')

# Earlier blocks on this page that this one builds on, re-run quietly first.
try:
    import numpy as np
    import numpy as np, time
except Exception:
    pass
_pm_earlier = [
    'import numpy as np\n\na = np.arange(12, dtype=np.int32).reshape(3, 4)\nprint(a)\n# [[ 0  1  2  3]\n#  [ 4  5  6  7]\n#  [ 8  9 10 11]]\n\nprint(a.dtype)         # int32\nprint(a.shape)         # (3, 4)\nprint(a.strides)       # (16, 4)   16 bytes to the next row, 4 bytes to the next column\nprint(a.nbytes)        # 48        12 ints * 4 bytes each',
    'import numpy as np\n\nnp.zeros((3, 4), dtype=np.float32)      # filled with 0.0, 12 floats = 48 bytes\nnp.ones((2, 2), dtype=np.int64)         # filled with 1\nnp.full((2, 3), 7.5)                    # filled with 7.5\nnp.empty((3, 3))                        # memory is not cleared, so overwrite every value before use\nnp.eye(4)                               # 4x4 identity matrix\nnp.arange(0, 10, 0.5)                   # like range, but allows a float step and returns an array\nnp.linspace(0, 1, 11)                   # 11 evenly spaced points from 0 to 1 inclusive',
    'rng = np.random.default_rng(seed=42)     # the modern API (NEP 19)\nrng.standard_normal((3, 4))              # 3x4 array from the standard normal N(0, 1)\nrng.uniform(0, 100, size=1000)           # 1000 uniform floats in [0, 100)\nrng.integers(0, 10, size=20)             # 20 ints in [0, 10)\nrng.choice(["A", "B", "C"], size=5, p=[0.5, 0.3, 0.2])',
    'import numpy as np\n\na = np.array([1, 2, 3, 4])\nb = np.array([10, 20, 30, 40])\n\na + b                       # [11, 22, 33, 44]   element by element\na * b                       # [10, 40, 90, 160]\na ** 2                      # [1, 4, 9, 16]\nnp.sqrt(a)                  # [1., 1.41, 1.73, 2.]\nnp.exp(a)                   # [2.72, 7.39, 20.09, 54.60]\na > 2                       # [False, False, True, True]   a boolean array\n\n# The "I want a Python loop" trap\nresult = np.empty_like(a)\nfor i in range(len(a)):\n    result[i] = a[i] ** 2 + b[i]      # SLOW: a Python loop over a C array\n\n# The vectorised way\nresult = a ** 2 + b                   # FAST: one pass in compiled C code',
    'import numpy as np\n\n# Subtract the column mean from each column (centre per feature)\nX = np.array([[1.0, 2.0, 3.0],\n              [4.0, 5.0, 6.0],\n              [7.0, 8.0, 9.0]])\n\ncol_mean = X.mean(axis=0)               # shape (3,)\ncentred  = X - col_mean                 # (3,3) - (3,) gives (3,3), and each column now has mean 0\n\n# Z-score normalisation per column\nstd = X.std(axis=0)\nz = (X - col_mean) / std\n\n# Outer product via broadcasting (no np.outer needed)\nu = np.array([1, 2, 3])                 # shape (3,)\nv = np.array([10, 20, 30, 40])          # shape (4,)\nouter = u[:, None] * v[None, :]         # (3,1) * (1,4) gives (3,4)\n# [[ 10  20  30  40]\n#  [ 20  40  60  80]\n#  [ 30  60  90 120]]\n\n# Pairwise distances (Euclidean) between rows of A and rows of B\nA = np.random.default_rng(0).standard_normal((100, 5))\nB = np.random.default_rng(1).standard_normal((50, 5))\n# (100,1,5) - (1,50,5) gives (100,50,5). Summing squares on the last axis gives (100,50)\ndists = np.sqrt(((A[:, None, :] - B[None, :, :]) ** 2).sum(axis=-1))\nprint(dists.shape)                      # (100, 50)',
    'import numpy as np\n\na = np.arange(20).reshape(5, 4)\n# [[ 0  1  2  3]\n#  [ 4  5  6  7]\n#  [ 8  9 10 11]\n#  [12 13 14 15]\n#  [16 17 18 19]]\n\na[[0, 2, 4]]                # rows 0, 2 and 4, shape (3, 4)\na[:, [3, 1, 0]]             # columns in a new order, shape (5, 3)\na[[0, 1, 2], [3, 2, 1]]     # picks (0,3), (1,2) and (2,1), giving [3, 6, 9]',
    'mask = a > 10\na[mask]                     # 1-D array of every element > 10: [11, 12, 13, 14, 15, 16, 17, 18, 19]\na[a % 2 == 0]               # all even values\n\n# Mask assignment: set every negative entry to zero\nx = np.array([-3, 1, -1, 5, -2])\nx[x < 0] = 0\nprint(x)                    # [0 1 0 5 0]',
    'x = np.array([-3, 1, -1, 5, -2])\nnp.where(x < 0, 0, x)       # [0, 1, 0, 5, 0]   the same result as np.maximum(x, 0)\nnp.where(x < 0, "neg", "pos")\n# array([\'neg\', \'pos\', \'neg\', \'pos\', \'neg\'], dtype=\'<U3\')\n\n# The one-argument form returns the indices where the condition is True\nidx = np.where(x < 0)\nprint(idx)                  # (array([0, 2, 4]),)\nprint(x[idx])               # [-3 -1 -2]',
    'import numpy as np\n\na = np.array([[1, 2, 3],\n              [4, 5, 6]])     # shape (2, 3)\n\na.sum()                       # 21    adds up every element\na.sum(axis=0)                 # [5, 7, 9]   collapses the rows, giving one sum per column\na.sum(axis=1)                 # [6, 15]     collapses the columns, giving one sum per row\n\na.mean(axis=0)                # [2.5, 3.5, 4.5]   the column means\na.argmax(axis=1)              # [2, 2]            the index of the max in each row',
    'import numpy as np\n\nA = np.array([[1, 2], [3, 4]], dtype=float)\nb = np.array([5, 6], dtype=float)\n\nA @ b                           # [17., 39.]    matrix times vector\nA @ A                           # matrix times matrix, a 2x2 result\nA.T @ A                         # Gram matrix\n\n# Solve Ax = b. Do NOT compute inv(A) @ b\nx = np.linalg.solve(A, b)       # [-4., 4.5]    faster and more accurate than the inverse\n\n# When you really need the inverse\nA_inv = np.linalg.inv(A)\n\n# Norms: Euclidean, Manhattan and Frobenius\nnp.linalg.norm(b)               # 7.81   L2 (Euclidean) by default\nnp.linalg.norm(b, ord=1)        # 11.0   L1 (Manhattan)\nnp.linalg.norm(A, ord="fro")    # 5.48   Frobenius\n\n# Eigendecomposition\neigvals, eigvecs = np.linalg.eig(A)\n# eigvals: [-0.37,  5.37]\n# eigvecs: columns are the corresponding unit eigenvectors\n\n# Singular value decomposition (the method behind PCA)\nU, s, Vt = np.linalg.svd(A, full_matrices=False)',
    'import numpy as np\n\ndata = np.array([2, 4, 4, 4, 5, 5, 7, 9])\n\nnp.percentile(data, [25, 50, 75])       # [4., 4.5, 5.5]   the quartiles\nnp.median(data)                         # 4.5\nnp.cov(data)                            # 4.57   sample covariance (for 1-D data, the variance)\n\n# Correlation matrix\nx = np.random.default_rng(0).standard_normal((100,))\ny = 2 * x + 0.5 * np.random.default_rng(1).standard_normal((100,))\nnp.corrcoef(x, y)\n# [[1.   0.97]\n#  [0.97 1.  ]]\n\n# Statistics on 2-D data: choose the axis carefully\nX = np.random.default_rng(0).standard_normal((1000, 5))   # 1000 samples, 5 features\nfeature_mean = X.mean(axis=0)           # shape (5,)   one mean per feature\nfeature_std  = X.std(axis=0, ddof=1)    # sample std (N-1 denominator)',
    'import numpy as np\n\na = np.arange(12)\na.reshape(3, 4)                 # 3x4 view\na.reshape(2, -1)                # -1 means "work it out", so 2x6\na.reshape(2, 2, 3)              # 3-D view\n\n# Transpose: swap or reorder the axes\nM = np.arange(24).reshape(2, 3, 4)\nM.T.shape                       # (4, 3, 2)   reverses all axes\nM.transpose(1, 0, 2).shape      # (3, 2, 4)   axes in the order you give\n\n# Stacking and concatenation\na = np.array([1, 2, 3])\nb = np.array([4, 5, 6])\n\nnp.concatenate([a, b])          # [1, 2, 3, 4, 5, 6]\nnp.stack([a, b])                # [[1,2,3],[4,5,6]]   adds a new axis at the front\nnp.vstack([a, b])               # [[1,2,3],[4,5,6]]   stacks them as rows\nnp.hstack([a, b])               # [1, 2, 3, 4, 5, 6]  joins 1-D arrays end to end\nnp.column_stack([a, b])         # [[1,4],[2,5],[3,6]] (each array becomes a column)\n\n# Splitting: the opposite of stacking\nM = np.arange(12).reshape(3, 4)\nnp.split(M, 2, axis=1)          # split into 2 along the columns, giving a list of (3,2) arrays',
    'import numpy as np\n\na = np.arange(10)\nb = a[2:6]                      # VIEW: shares memory with a\nb[0] = 999\nprint(a)                        # [  0   1 999   3   4   5   6   7   8   9]   a changed too!\n\nc = a[[2, 3, 4]]                # COPY: fancy indexing always copies\nc[0] = -1\nprint(a)                        # unchanged\n\n# Check explicitly\nprint(b.base is a)              # True   b is a view of a\nprint(c.base is a)              # False  c is a separate copy\n\n# Force a copy when you want one\nd = a[2:6].copy()\nd[0] = 555\nprint(a)                        # unchanged',
    'import numpy as np, time\n\nn = 1_000_000\na = np.random.default_rng(0).standard_normal(n)\nb = np.random.default_rng(1).standard_normal(n)\n\n# Python loop\nt0 = time.perf_counter()\nresult_loop = [a[i] * b[i] + a[i] for i in range(n)]\nt_loop = time.perf_counter() - t0\n\n# NumPy vectorised\nt0 = time.perf_counter()\nresult_vec = a * b + a\nt_vec = time.perf_counter() - t0\n\nprint(f"loop: {t_loop:.3f}s   vec: {t_vec:.4f}s   speedup: {t_loop/t_vec:.0f}x")\n# loop: ~0.300s   vec: ~0.003s   speedup: ~100x',
    'import numpy as np\n\nA = np.random.default_rng(0).standard_normal((100, 50))\nB = np.random.default_rng(1).standard_normal((50, 30))\n\n# Standard matmul\nnp.einsum("ij,jk->ik", A, B).shape          # (100, 30)   the same as A @ B\n\n# Batched dot product of corresponding rows\nX = np.random.default_rng(0).standard_normal((1000, 5))\nY = np.random.default_rng(1).standard_normal((1000, 5))\ndots = np.einsum("ij,ij->i", X, Y)          # shape (1000,) (one dot product per row)',
    'import numpy as np\n\na = np.arange(12).reshape(3, 4)\n\nnp.save("a.npy", a)                         # binary, fastest, and keeps the dtype\nb = np.load("a.npy")\n\nnp.savez("bundle.npz", x=a, y=a * 2)        # multiple arrays in one file\ndata = np.load("bundle.npz")\nprint(data["x"], data["y"])\n\nnp.savetxt("a.csv", a, delimiter=",", fmt="%d")   # human-readable, but can lose float precision\nb = np.loadtxt("a.csv", delimiter=",")',
    'import numpy as np\n\na = np.array([100, 100], dtype=np.int8)     # range is [-128, 127], so 100 fits\nprint(a + a)                                # [-56 -56]  it WRAPPED around with no warning',
]
def _pm_run_earlier(_srcs, _g):
    import sys as _s, io as _io, warnings as _w
    for _src in _srcs:
        _o, _e = _s.stdout, _s.stderr
        _s.stdout = _s.stderr = _io.StringIO()
        try:
            with _w.catch_warnings():
                _w.simplefilter('ignore')
                exec(compile(_src, '<earlier block>', 'exec'), _g)
        except (Exception, SystemExit):
            pass
        finally:
            _s.stdout, _s.stderr = _o, _e
    if 'matplotlib.pyplot' in _s.modules:
        _s.modules['matplotlib.pyplot'].close('all')   # their plots belong to their own blocks
_pm_run_earlier(_pm_earlier, globals())
del _pm_earlier, _pm_run_earlier

When NumPy can't find a common type, it uses a string dtype or object (an array of pointers to Python objects). You lose the speed of vectorisation, and an object array is no faster than a list. Replace missing values with np.nan (which is a float) before you create the array, or use pandas for data with mixed types.

3. In-place ops on shared views

python
import numpy as np

a = np.arange(10)
b = a[::2]                  # view of every other element: [0, 2, 4, 6, 8]
b *= 100                    # IN-PLACE, so this also changes a
print(a)                    # [  0   1 200   3 400   5 600   7 800   9]

*=, += and similar operators change the underlying buffer. If b is a view, a sees the change too. If you don't want that, use b = b * 100 (which creates a new array), or call .copy() before the in-place operation.

4. Using Python functions when NumPy ops exist

python
import math
import numpy as np

x = np.linspace(0, 10, 1_000_000)

# Slow: a Python loop hidden inside a list comprehension
y = [math.sin(v) for v in x]                # ~200ms

# Fast: a vectorised loop in C
y = np.sin(x)                               # ~5ms

math.sin, math.exp, math.log are scalar functions. np.sin, np.exp, np.log are vectorised. Using the math versions on arrays loses the benefit of NumPy. The same applies to min/max: np.min(a) is much faster than min(a) on a large array.

5. Forgetting keepdims and breaking broadcast

python
import numpy as np

a = np.random.default_rng(0).standard_normal((100, 5))

# Wrong: the shapes do not match
centred = a - a.mean(axis=1)                # (100, 5) - (100,) fails because the shapes do not line up

# Right: keep the collapsed axis with length 1
centred = a - a.mean(axis=1, keepdims=True) # (100, 5) - (100, 1) broadcasts correctly

Whenever you reduce an array and then combine the result with the original, you almost always want keepdims=True.


🎯 Your Turn: Rolling Mean in O(n)

Write moving_average(arr, window) that returns the rolling mean of arr with a fixed window size. The simple solution is O(n × window). The better one is O(n), using np.cumsum.

Expected behaviour:

python
moving_average(np.array([1, 2, 3, 4, 5, 6, 7]), window=3)
# array([2., 3., 4., 5., 6.])           means of [1,2,3], [2,3,4], ..., [5,6,7]

moving_average(np.array([1.0, 2.0, 4.0, 8.0]), window=2)
# array([1.5, 3.0, 6.0])
+ setup added so this can run · defines moving_average, np
# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

def moving_average(*_a, **_kw):
    print('-> moving_average() called')
    return _AutoMock('moving_average()')
np = _AutoMock('np')

Constraints:

  • Return a 1-D ndarray of length len(arr) - window + 1.
  • Must run in O(n), in a single pass, using np.cumsum, not a Python loop.
  • Raise ValueError if window < 1 or window > len(arr).

Skeleton:

python
import numpy as np

def moving_average(arr, window):
    # TODO 1: validate window (>= 1, <= len(arr))
    # TODO 2: compute cumulative sum, prepended with a 0
    # TODO 3: subtract cumsum[:-window] from cumsum[window:] to get window sums
    # TODO 4: divide by window to get means
    ...
Hint 1: Cumsum gives you prefix sums np.cumsum(arr) returns [arr[0], arr[0]+arr[1], arr[0]+arr[1]+arr[2], ...]. The sum of any window arr[i:i+w] equals cumsum[i+w-1] - cumsum[i-1]. Prepend a zero (np.concatenate([[0], cumsum])) so the i=0 case works without a special branch.
Hint 2: Vectorise the subtraction With the prepended cumsum (length n+1), the array of window sums is cumsum[window:] - cumsum[:-window]. Both are length n - window + 1. Divide by window (a scalar, broadcasts automatically) to get means. No loop.
Show full solution
python
import numpy as np

def moving_average(arr, window):
    """Rolling mean of `arr` over a fixed-size window. Runs in O(n) time using cumsum."""
    if window < 1:
        raise ValueError("window must be >= 1")
    if window > len(arr):
        raise ValueError("window cannot exceed array length")

    cumsum = np.cumsum(arr, dtype=np.float64)            # prefix sums
    cumsum = np.concatenate(([0.0], cumsum))             # put a 0 at the front so cumsum[0] == 0
    window_sums = cumsum[window:] - cumsum[:-window]     # the sum of each window
    return window_sums / window


# Verify
print(moving_average(np.array([1, 2, 3, 4, 5, 6, 7]), 3))
# [2. 3. 4. 5. 6.]

print(moving_average(np.array([1.0, 2.0, 4.0, 8.0]), 2))
# [1.5 3.  6. ]

# Speed: the naive loop compared with cumsum on a million elements
import time
arr = np.random.default_rng(0).standard_normal(1_000_000)
w = 100

t0 = time.perf_counter()
naive = np.array([arr[i:i+w].mean() for i in range(len(arr) - w + 1)])
t_naive = time.perf_counter() - t0

t0 = time.perf_counter()
fast = moving_average(arr, w)
t_fast = time.perf_counter() - t0

print(f"naive: {t_naive:.3f}s   cumsum: {t_fast:.4f}s   speedup: {t_naive/t_fast:.0f}x")
# naive: ~1.2s   cumsum: ~0.01s   speedup: ~100x

Why this is O(n): cumsum is a single linear pass. The subtraction and division are also linear, and they run vectorised in C. There is no nested loop anywhere. Doubling the window size doesn't change the runtime.

Why dtype=np.float64 on cumsum: if arr has an integer dtype, the cumulative sums can overflow without warning. Converting to float64 first avoids that. For very long arrays of large floats, even float64 collects rounding error, because np.cumsum adds the numbers in the simplest way. Rolling means in production code use Kahan summation or the algorithm in pandas' rolling().mean(), which subtracts the leaving element rather than recomputing from a global prefix sum (better numerics, same O(n)).

This technique, prefix sums for sliding-window aggregates, works for many problems. Rolling sum, rolling mean, "is there a window of size k with a sum of at least X?" and "longest subarray with sum < S" can all be solved with operations on a cumsum array. It is worth practising until you know it well.


What You Learned

  • An ndarray is a contiguous buffer + dtype + shape + strides. Reshaping and slicing change the metadata, not the data.
  • Use the right constructor (zeros, empty, full, eye, linspace) and the modern default_rng for random data.
  • Vectorise: write computations as whole-array operations, never per-element Python loops. Expect 50–200× speedups.
  • Broadcasting aligns shapes from the right, and size-1 and missing dimensions stretch. Use [:, None] and keepdims=True to control alignment.
  • Advanced indexing: integer arrays and boolean masks both copy, but plain slices are views. np.where for vectorised ternaries.
  • axis=k collapses axis k. axis=0 gives per-column results, axis=1 per-row.
  • Linear algebra: prefer @ over np.dot, np.linalg.solve(A, b) over inv(A) @ b.
  • Views vs copies: slices are views, and fancy indexing copies. Changes to a view also change the original. Use .copy() when you are not sure.
  • Watch for integer overflow without warning, accidental object dtype, and the in-place-on-shared-view trap.
  • For loops that can't be vectorised, use numba or Cython. For tensor contractions, use np.einsum.

Next: Pandas DataFrames: Command Center: the same arrays, with labels, mixed column types and grouping.

Practice this

on practicepython.in

Short exercises that run in your browser and tell you what your code did, as well as whether a test passed.

Found this useful? Share it with a friend who’s learning Python.