NumPy Arrays, Indexing, and Broadcasting

You have a list of prices and you want to apply a 10% increase. The obvious thing to type is this:

prices = [10, 20, 30]print(prices * 2)      # [10, 20, 30, 10, 20, 30]

Python duplicated the list instead of doubling the numbers, because * on a list means "repeat". So you write the loop instead, and it works — until the list has fifty million entries and you are staring at a progress bar.

1import timeit2import numpy as np34data = list(range(5_000_000))5arr = np.arange(5_000_000)67# best of 5 runs, so one slow run does not distort the comparison8t_list = min(timeit.repeat(lambda: [x * 1.1 for x in data], number=1, repeat=5))9t_numpy = min(timeit.repeat(lambda: arr * 1.1, number=1, repeat=5))1011print(f"pure Python: {t_list:.3f}s")     # ~0.2s on a recent laptop12print(f"NumPy:       {t_numpy:.4f}s")    # ~0.003s

Tens of times faster — about fifty times on a recent laptop, though the exact ratio depends on your machine — and the code says what it means. That gap is not a clever optimisation. It comes from a difference in how the two things are stored in memory, and understanding that difference explains almost everything else about NumPy.

Why a list is slow and an array is not

A Python list is an array of pointers. Each element points somewhere else in memory, to a full Python object that carries its own type tag and reference count. Adding two lists element-wise means: follow a pointer, check the type, unbox the value, add, allocate a new object, store a pointer. Per element.

A NumPy array is one unbroken block of memory holding raw numbers of a single type, with the type recorded once for the whole array.

Python list [1, 2, 3]        NumPy array [1, 2, 3]┌───┬───┬───┐                ┌────┬────┬────┐│ • │ • │ • │  pointers      │ 1  │ 2  │ 3  │  dtype=int64, recorded once└─┼─┴─┼─┴─┼─┘                └────┴────┴────┘  ▼   ▼   ▼                  24 contiguous bytes int int int  (28 bytes each, scattered)

Because the block is contiguous and uniformly typed, the addition loop can run in compiled C, use the CPU's vector instructions, and never touch the Python interpreter. That is where the speed-up comes from.

The price of that speed is rigidity: an array holds exactly one type, and its size is fixed at creation. You trade flexibility for a hundredfold reduction in per-element overhead.

Creating arrays

1import numpy as np23a = np.array([1, 2, 3])                    # from a list4b = np.array([[1, 2, 3], [4, 5, 6]])       # 2D from nested lists56np.zeros((2, 3))          # 2x3 of 0.07np.ones(5)                # [1. 1. 1. 1. 1.]8np.full((2, 2), 7)        # filled with 79np.eye(3)                 # identity matrix10np.arange(0, 10, 2)       # [0 2 4 6 8]  -- stop is excluded11np.linspace(0, 1, 5)      # [0. 0.25 0.5 0.75 1.]  -- stop IS included1213rng = np.random.default_rng(seed=42)14rng.normal(loc=0, scale=1, size=(3, 3))    # reproducible random numbers

arange and linspace disagree about the endpoint, which catches everyone once. arange takes a step and excludes the stop; linspace takes a count and includes it. For plotting axes and grid searches you almost always want linspace.

Every array knows four things about itself

1b = np.array([[1, 2, 3], [4, 5, 6]])2print(b.shape)    # (2, 3)   -- rows, columns3print(b.ndim)     # 2        -- number of dimensions4print(b.size)     # 6        -- total elements5print(b.dtype)    # int64    -- the single type of every element

shape is the one you will check constantly. Nearly every NumPy error message is really a shape mismatch wearing a disguise.

Dtypes, and the overflow that does not warn you

Unlike Python integers, NumPy integers have a fixed width and they wrap around silently:

x = np.array([127], dtype=np.int8)print(x + 1)      # [-128]  -- no error, no warning, completely wrong

This is a real hazard when you shrink dtypes to save memory on a large dataset. Converting a pixel count column to int8 to save space works fine until a value exceeds 127, at which point your data becomes negative and every downstream statistic is poisoned. Halving memory by moving from float64 to float32 is usually safe for machine learning; shrinking integers needs you to check the maximum first.

Indexing and slicing

One dimension behaves like a list

1a = np.array([10, 20, 30, 40, 50])2print(a[0], a[-1])     # 10 503print(a[1:4])          # [20 30 40]4print(a[::2])          # [10 30 50]5print(a[::-1])         # [50 40 30 20 10]

Two dimensions use one bracket, not two

1m = np.array([[1, 2, 3],2              [4, 5, 6],3              [7, 8, 9]])45print(m[1, 2])       # 6   -- row 1, column 26print(m[1])          # [4 5 6]      whole row7print(m[:, 1])       # [2 5 8]      whole column8print(m[0:2, 1:3])   # [[2 3]9                     #  [5 6]]      a sub-block

m[1][2] also works but is slower — it builds an intermediate row array and then indexes that. Use the comma form.

Boolean indexing: the workhorse

Comparing an array to a value gives an array of booleans, and an array of booleans can be used as a filter. This is how essentially all data filtering happens.

1temps = np.array([18.5, 22.1, 31.7, 15.2, 28.9])23mask = temps > 254print(mask)             # [False False  True False  True]5print(temps[mask])      # [31.7 28.9]6print(mask.sum())       # 2  -- True counts as 1, so this counts matches7print(temps[temps > 25].mean().round(2))   # 30.389# combine conditions with & and | -- and parenthesise EVERY term10print(temps[(temps > 20) & (temps < 30)])    # [22.1 28.9]

Two rules that trip people up. You must use & and |, not and and or — the Python keywords try to evaluate the whole array as a single true/false and raise ValueError: The truth value of an array with more than one element is ambiguous. And you must wrap each comparison in brackets, because & binds more tightly than >, so leaving them off silently compares the wrong things.

temps[temps > 25] = 25         # assignment through a mask: cap the outliersprint(np.where(temps > 20, "warm", "cold"))   # element-wise if/else

Fancy indexing

1a = np.array([10, 20, 30, 40, 50])2print(a[[0, 2, 4]])       # [10 30 50]  -- pick by position list3print(a[[2, 2, 0]])       # [30 30 10]  -- repeats allowed, order is yours45order = np.argsort(a)[::-1]6print(a[order])           # [50 40 30 20 10]  -- sort descending

Views and copies: the trap that corrupts your data

A slice of a Python list is a copy. A slice of a NumPy array is a view — a window onto the same memory. Writing through the window changes the original.

1original = np.array([1, 2, 3, 4, 5])2window = original[1:4]3window[0] = 9994print(original)      # [  1 999   3   4   5]  -- the original changed

This is deliberate and it is what makes slicing free — no data is moved. But it means the innocent-looking line train = data[:800] followed by any in-place edit of train quietly rewrites your source data. Take an explicit copy when you intend independence:

window = original[1:4].copy()print(window.base is None)     # True -- this array owns its own memory

The rule is learnable: basic slicing gives a view; boolean masks and fancy indexing give copies. When in doubt, check arr.base — if it is not None, you are holding a view.

Slicing an array does not copy it. If you plan to modify a slice, call .copy() first, or you will change data you did not mean to touch.

Vectorised operations

Arithmetic on arrays applies element by element, with no loop written and no loop run in Python:

1a = np.array([1, 2, 3, 4])2b = np.array([10, 20, 30, 40])34print(a + b)        # [11 22 33 44]5print(b / a)        # [10. 10. 10. 10.]6print(a ** 2)       # [ 1  4  9 16]7print(np.sqrt(b))   # [3.16227766 4.47213595 5.47722558 6.32455532]8print(np.log(b))    # element-wise natural log

Aggregation, and the axis argument everyone gets backwards

1m = np.array([[1, 2, 3],2              [4, 5, 6]])34print(m.sum())          # 21  -- everything5print(m.sum(axis=0))    # [5 7 9]    -- collapses rows,    one result per COLUMN6print(m.sum(axis=1))    # [ 6 15]    -- collapses columns, one result per ROW

The mental model that sticks: axis names the dimension that disappears. The array is (2, 3); axis=0 removes the first dimension and leaves (3,). So to get a mean per feature, where features are columns, you want axis=0. Getting this wrong gives you a result of the wrong length, which usually raises an error later rather than at the point of the mistake.

The same axis logic applies to mean, std, min, max, argmin, argmax, median and percentile. And use the nan-prefixed versions when missing values are present, because a single NaN makes an ordinary mean return NaN for the whole thing:

vals = np.array([1.0, 2.0, np.nan, 4.0])print(vals.mean())        # nanprint(np.nanmean(vals))   # 2.3333333333333335

Broadcasting

You have a data matrix of 1,000 samples by 4 features, and you want to subtract each feature's mean. The means are 4 numbers; the data is 4,000 numbers. Broadcasting is the rule that lets you write data - means and have NumPy work out that the 4 means should be reused down every row.

1data = np.array([[10., 20., 30.],2                 [40., 50., 60.]])          # shape (2, 3)34col_means = data.mean(axis=0)               # shape (3,)  -> [25. 35. 45.]5centred = data - col_means6print(centred)     # [[-15. -15. -15.]7                   #  [ 15.  15.  15.]]

The rule, exactly

Line the shapes up from the right. Two dimensions are compatible if they are equal, or if one of them is 1. A missing leading dimension is treated as 1.

data       (2, 3)col_means     (3,)   ->  treated as (1, 3)                          ^ stretched down to 2 rowsresult     (2, 3)   OKdata       (2, 3)other         (2,)   ->  treated as (1, 2)                          3 vs 2 -- neither is 1  -> ValueError

That last row is the everyday failure. You want to subtract a per-row value from a (3, 4) matrix, and your row values have shape (3,). Reshape them into a column first:

row_means = data.mean(axis=1)              # shape (2,)  -- wrong orientationcentred = data - row_means[:, np.newaxis]  # shape (2, 1) -- now it broadcasts

[:, np.newaxis] (equivalently .reshape(-1, 1)) turns a flat array into a single column. Reach for it whenever an operation "should" work but the shapes disagree.

Reshaping and combining

1a = np.arange(12)2print(a.reshape(3, 4).shape)      # (3, 4)3print(a.reshape(3, -1).shape)     # (3, 4)  -- -1 means "work it out"4print(a.reshape(3, 4).T.shape)    # (4, 3)  -- transpose5print(a.reshape(3, 4).ravel())    # flat again67x = np.array([[1, 2], [3, 4]])8y = np.array([[5, 6]])910print(np.vstack([x, y]).shape)          # (3, 2)  stack rows11print(np.hstack([x, y.T]).shape)        # (2, 3)  stack columns12print(np.concatenate([x, y], axis=0))   # the general form

reshape returns a view where it can, so it is nearly free; it fails only when the total element count does not match, which is a helpful error rather than a silent one.

What this looks like in real feature preparation

Standardising features — putting every column on a mean of 0 and a standard deviation of 1 — is the single most common preprocessing step before any distance-based or gradient-based model. The formula is $z = \frac{x - \mu}{\sigma}$, computed per column, and broadcasting makes it three lines:

1import numpy as np23rng = np.random.default_rng(0)4X = np.column_stack([5    rng.normal(50, 10, 1000),        # feature 1: around 506    rng.normal(0.5, 0.1, 1000),      # feature 2: around 0.57    rng.normal(20000, 5000, 1000),   # feature 3: around 200008])910mu = X.mean(axis=0)                  # (3,)  one mean per column11sigma = X.std(axis=0)                # (3,)12Z = (X - mu) / sigma                 # broadcasts down all 1000 rows1314np.set_printoptions(suppress=True)  # plain decimals, not 1.9e+0415print(X.mean(axis=0).round(1))       # [   49.5     0.5 19770.7]16print(Z.mean(axis=0).round(6))       # [-0. -0.  0.]  (-0. is just zero)17print(Z.std(axis=0).round(6))        # [1. 1. 1.]

Without scaling, the third feature's values are a thousand times larger than the second's, so any model that measures distance would treat feature 2 as if it did not exist. With three lines of broadcasting, all three carry equal weight.

Two habits make NumPy code reliable. First, print .shape whenever an operation surprises you — the answer is a shape mismatch far more often than a logic error. Second, if you find yourself writing a for loop over array elements, stop and look for the vectorised form. It exists almost always, it is shorter, and it is the difference between a script that finishes and one you abandon.