Thinking in arrays
By the end: You can manipulate data without loops, which is how every ML library expects it.
- Vectorised operations
- Indexing, masking and reshaping
- pandas for tabular work
- Group-by and joins
New to numpy & pandas? NumPy is a Python library for holding large grids of numbers and doing arithmetic on all of them in one instruction, rather than one at a time in a loop. pandas builds on it for labelled tables with named columns and mixed types, which is the shape real data arrives in; its table is called a DataFrame, and you filter, join and group it in memory on your own machine. Why it matters
Every numerical library in this field expects whole-array operations rather than loops, and the gap is not small: the loop version of a job can be a hundred times slower than the same thing expressed as one expression over the whole array. This section is about thinking that way, along with the indexing, masking and reshaping that make it possible, and the table operations in pandas that most real data work actually consists of.
Stop writing loops#
A loop over a NumPy array is usually a hundred times slower than the vectorised equivalent, because the loop runs in Python while the vectorised version runs in compiled code. Learning to see operations as whole-array is the main skill here.
import numpy as np
# Slow, and how most people first write it
result = []
for x in data:
result.append(x * 2 + 1)
# Fast, and shorter
result = data * 2 + 1Two things make this worth the adjustment. The obvious one is speed: the Python version pays interpreter overhead per element, boxing each number into an object, while the array version dispatches once into compiled code that walks contiguous memory. The less obvious one is that every library you will hand this data to — scikit-learn, PyTorch, anything — expects arrays, so a list of Python floats has to be converted anyway.
The axis argument#
This is the single thing people get wrong most often, and the rule is short: axis names the dimension that disappears. An array of shape (2, 3) summed along axis 0 loses the 2 and leaves shape (3,); summed along axis 1 it loses the 3 and leaves shape (2,). Say it that way rather than in terms of rows and columns, because the rows-and-columns version stops working the moment there are three dimensions.
a = np.array([[1, 2, 3],
[4, 5, 6]]) # shape (2, 3)
a.sum() # 21 everything, shape ()
a.sum(axis=0) # [5 7 9] shape (3,) the 2 is gone
a.sum(axis=1) # [ 6 15] shape (2,) the 3 is gone
a.sum(axis=1, keepdims=True).shape # (2, 1) kept as a length-1 dimension
a.argmax(axis=1) # [2 2] index of the max in each rowThe same rule applies to any array. For a batch of token embeddings of shape (batch, sequence, features), averaging over axis=1 collapses the sequence and gives one vector per item, which is mean pooling. Averaging over axis=0 collapses the batch, which is almost never what anyone meant. Negative indices count from the end, so axis=-1 is the last dimension and is the safe way to write an operation that should apply to features regardless of how many leading dimensions there are.
def softmax(z):
"""Over the last axis, so it works on (n,), (batch, n) and (batch, seq, n)."""
z = z - z.max(axis=-1, keepdims=True) # subtract the max: stops exp overflowing
e = np.exp(z)
return e / e.sum(axis=-1, keepdims=True)
logits = np.array([[2.0, 1.0, 0.1],
[0.5, 0.5, 0.5]])
softmax(logits)
# [[0.659 0.242 0.099]
# [0.333 0.333 0.333]]
softmax(logits).sum(axis=-1) # [1. 1.] always check thisThe axis mistake that does not raise
Normalising rows with m / m.sum(axis=1) fails loudly on a (2, 3) array, because (2,3) and (2,) cannot broadcast. On a square array it succeeds and normalises the columns instead. Take m = [[1, 1], [1, 3]]: the correct row-normalised result is [[0.5, 0.5], [0.25, 0.75]], and the version without keepdims gives [[0.5, 0.25], [0.5, 0.75]], whose rows sum to 0.75 and 1.25. Nothing errors, the numbers look like probabilities, and a model trains on them. Write keepdims=True, and assert the rows sum to one.
Broadcasting, as rules#
Broadcasting is how NumPy lets arrays of different shapes take part in the same operation without you writing a loop or copying data. It is worth learning as three mechanical rules rather than as intuition, because then you can work out what a line does without running it.
1. Line the shapes up from the RIGHT, padding the shorter one with 1s on the left.
2. Two dimensions are compatible if they are equal, or if one of them is 1.
3. A dimension of 1 is stretched to match the other. Anything else is an error.
(3, 2) (3, 1) (2, 3) (3, 2)
+ ( 2) + ( 2) + (2, 1) + ( 3)
--------- --------- -------- ---------
(3, 2) (3, 2) (2, 3) error: 2 and 3 are neither
ok ok ok equal nor 1x = np.arange(6).reshape(3, 2) # [[0 1] [2 3] [4 5]]
row = np.array([10, 20]) # (2,) -> treated as (1, 2), stretched down rows
col = np.array([[100], [200], [300]]) # (3, 1) -> stretched across columns
x + row
# [[10 21]
# [12 23]
# [14 25]]
x + col
# [[100 101]
# [202 203]
# [304 305]]
x + np.array([1, 2, 3])
# ValueError: operands could not be broadcast together with shapes (3,2) (3,)Nothing is copied. A stretched dimension is read repeatedly from the same memory, which is why subtracting a mean vector from a million rows costs nothing extra in space. The tool for making shapes line up is None, also spelled np.newaxis, which inserts a dimension of length 1 wherever you put it.
v = np.array([1, 2, 3])
v[:, None].shape # (3, 1) a column
v[None, :].shape # (1, 3) a row
# All pairwise differences, with no loop:
v[:, None] - v[None, :]
# [[ 0 -1 -2]
# [ 1 0 -1]
# [ 2 1 0]]
# Centre every feature of a (n, d) matrix on its own mean:
X - X.mean(axis=0) # (n,d) - (d,) -> fine, stretches down rows
# Scale every row to unit length:
X / np.linalg.norm(X, axis=1, keepdims=True) # (n,d) / (n,1)Masking#
A boolean array selects elements. This replaces most filtering loops and composes cleanly.
scores = np.array([0.2, 0.9, 0.5, 0.95])
scores > 0.8 # array([False, True, False, True])
scores[scores > 0.8] # array([0.9, 0.95])
np.where(scores > 0.8, 1, 0) # array([0, 1, 0, 1])
(scores > 0.4) & (scores < 0.9) # & and |, not and/orThree things follow from this. A mask is itself an array, so masks combine with & and | and invert with ~, and each operand needs its own brackets because & binds tighter than the comparison operators. A mask counts as well as selects, because summing booleans counts the Trues. And indexing with a mask, or with a list of indices, returns a copy rather than a view, which is the opposite of what a plain slice does.
mask = (scores > 0.4) & (scores < 0.9)
mask.sum() # 1 how many match
mask.mean() # 0.25 what fraction match
scores[~mask] # everything that does not
np.where(mask)[0] # array([2]) the positions, not the values
np.shares_memory(scores, scores[1:3]) # True: a slice is a view
np.shares_memory(scores, scores[mask]) # False: a mask gives a copy
np.shares_memory(scores, scores[[0, 2]]) # False: so does a list of indices
scores[scores > 0.8] = 0.8 # assignment through a mask does write backShapes#
a = np.arange(12)
a.reshape(3, 4) # explicit
a.reshape(3, -1) # -1 means "work it out"
a[:, None].shape # (12, 1) add an axis
a.T # transpose
np.concatenate([a, a]) # join along an existing axis
np.stack([a, a]) # join along a new axisreshape never moves data; it reinterprets the same buffer, so it is free and it returns a view. It fills the new shape in row-major order, last axis fastest, which is why reshaping (2, 3) to (3, 2) is not a transpose and gives a different arrangement of the same numbers. Transposing also returns a view, one whose strides are swapped rather than whose memory is rearranged, which is why some later operation on it may quietly copy.
a = np.arange(6).reshape(2, 3)
# [[0 1 2]
# [3 4 5]]
a.reshape(3, 2) # [[0 1] [2 3] [4 5]] filled across, then down
a.T # [[0 3] [1 4] [2 5]] not the same thing
np.concatenate([a, a], axis=0).shape # (4, 3) more rows
np.concatenate([a, a], axis=1).shape # (2, 6) more columns
np.stack([a, a]).shape # (2, 2, 3) a new leading axis
a.squeeze() # drop every axis of length 1
a.ravel() # flatten to 1-D, a view when it can bePrint the shape, not the array
Almost all debugging in array code is shape debugging, and the array is too big to read. Put print(name, x.shape) at each step, or assert x.shape == (batch, n_classes) at the boundaries of a function. An assert on shape catches the class of bug where the numbers are wrong but plausible, which no amount of looking at the output will.
pandas#
A DataFrame is a table with labelled columns. The operations that matter are selection, grouping and joining, which cover most data preparation.
import pandas as pd
df = pd.read_csv("orders.csv")
df.head() # look at it, always
df.info() # types and null counts
df.describe() # distributions
df[df.total > 100] # filter rows
df[["id", "total"]] # select columns
df.loc[df.total > 100, "flag"] = True # conditional assignment
df.groupby("customer_id").agg(
orders=("id", "count"),
revenue=("total", "sum"),
)
pd.merge(orders, customers, on="customer_id", how="left")Two pieces of vocabulary before anything else. A DataFrame has an index as well as columns, and the index is what alignment happens on: adding two Series lines them up by index label, not by position, so a subtraction between a filtered frame and an unfiltered one produces NaN wherever the labels do not match rather than the answer you expected. And there are two accessors: .loc selects by label and .iloc by position, and mixing them up on a frame whose index is not 0..n-1 gives wrong rows silently.
df.loc[3, "total"] # the row whose index LABEL is 3
df.iloc[3] # the FOURTH row, whatever its label
df.loc[df.total > 100, "flag"] = True # label-based, and the way to assign
sub = df[df.total > 100] # the index keeps the original labels: 4, 9, 17...
sub.reset_index(drop=True) # renumber when position is what you wantAssign with .loc, never through two brackets
Writing df[df.total > 100]["flag"] = True is chained assignment: the first bracket produces a new object, the assignment lands on that, and the original frame is unchanged. Older pandas raised SettingWithCopyWarning here, a warning people learned to ignore because it also fired on code that worked. Under the copy-on-write behaviour that is the default from pandas 3 it raises ChainedAssignmentError and tells you plainly that chained assignment never updates the original. Both are saying the same thing, and df.loc[df.total > 100, "flag"] = True is the fix in every version.
The related habit is to call .copy() when you deliberately want an independent frame to work on, as in ok = df[df.status == 'ok'].copy(). It states the intention, and it means later columns you add are yours rather than an attempt to write through a filter.
Group-by and merge, the two that matter#
Group-by splits the frame by a key, applies something to each piece, and combines the results. Which of the three shapes you want depends on what comes out: agg gives one row per group, transform gives one value per original row, and apply gives whatever your function returns and is the slow one to avoid when either of the others will do.
calls = pd.DataFrame({
"user_id": [1, 1, 2, 2, 3],
"model": ["sonnet", "haiku", "sonnet", "sonnet", "haiku"],
"tokens": [1200, 300, 900, 1500, None],
"latency_ms": [820, 140, 760, 2100, 155],
})
calls.groupby("model").agg(
calls=("user_id", "size"), # rows in the group
users=("user_id", "nunique"),
tokens=("tokens", "sum"),
p95=("latency_ms", lambda s: s.quantile(0.95)),
)
# calls users tokens p95
# model
# haiku 2 2 300.0 154.25
# sonnet 3 2 3600.0 1972.00Note haiku's tokens: two rows, one of them missing, and the sum is 300 rather than an error or a NaN. size counts rows in the group and count counts non-null values in a column, so on that group they give 2 and 1. Choosing the wrong one is how a denominator quietly becomes too small.
calls.groupby("model").size() # haiku 2, sonnet 3 rows
calls.groupby("model")["tokens"].count() # haiku 1, sonnet 3 non-null
# transform: one value per original row, so it can go straight back on the frame
calls["model_total"] = calls.groupby("model")["tokens"].transform("sum")
# user_id model tokens latency_ms model_total
# 0 1 sonnet 1200.0 820 3600.0
# 1 1 haiku 300.0 140 300.0
# ...
calls["share"] = calls["tokens"] / calls["model_total"]merge is a SQL join with a different spelling, and the same things go wrong. Two arguments turn its failure modes into errors instead of wrong numbers: validate states the relationship you believe holds and raises if it does not, and indicator adds a column saying which side each row came from so you can count the rows that failed to match.
users = pd.DataFrame({"user_id": [1, 2, 4], "plan": ["pro", "free", "pro"]})
m = calls.merge(users, on="user_id", how="left",
validate="many_to_one", indicator=True)
m["_merge"].value_counts()
# both 4
# left_only 1 <- user 3 has no row in users
# right_only 0
# If users had two rows for user_id 1, the same call raises rather than silently
# duplicating: MergeError: Merge keys are not unique in right dataset;
# not a many-to-one merge
#
# Without validate, 5 rows in become 7 rows out and every sum is now wrong.The habit worth forming is to print len(df) before and after every merge. A row count that changed when you expected a lookup is the single most common cause of a number being wrong in a data pipeline, and it is free to check.
Missing values#
df.isna().sum() # how many per column
df.dropna(subset=["total"]) # drop rows missing something essential
df.total.fillna(df.total.median()) # impute
# Whether a value is missing is often informative in itself
df["total_was_missing"] = df.total.isna()Where people go wrong
Fitting an imputer or scaler on the whole dataset before splitting. Information from the test set enters training, the score improves, and the model fails in production.
If this comes up in an interview
Why is vectorised NumPy faster than a Python loop?
The loop executes in the Python interpreter with per-element overhead, while the vectorised operation runs in compiled code over contiguous memory, often using multiple values per instruction. The difference is typically one to two orders of magnitude.
What does SettingWithCopyWarning mean?
You may be assigning to a copy rather than the original DataFrame, so the change silently does nothing. It usually follows a filter, and using .loc for conditional assignment avoids both the warning and the bug.
How would you handle missing values?
First find out whether missingness is informative, and if so record it as a feature before imputing. Then drop rows missing something essential, and impute the rest with a median or a model, fitted inside the training fold to avoid leakage.
Build this: Take a messy CSV and clean it without a single explicit loop: handle missing values, filter, group and join. Then time it against a loop-based version.
Watch it explained
This is one section of NumPy & pandas, 18 hours across 4 sections. Read it in full to get the whole topic in order.