NumPy Interview Questions for Machine Learning, with Code
NumPy interview questions for machine learning and data science, with short answers and working code for broadcasting, indexing, views and copies, einsum, stable softmax and top-k, plus practice problems you solve in code.
What NumPy questions are asked in machine learning interviews?
Machine learning and data science interviews ask NumPy questions about how arrays differ from lists, dtypes, shapes and reshaping, views and copies, indexing, broadcasting, the axis argument, vectorising loops, matrix multiplication, einsum, random numbers, numerical stability and common ML operations such as one-hot encoding, pairwise distances, normalisation and top-k. TensorTonic has a NumPy Sheet of 25 problems in five sections for these topics, and this guide gives a short answer and a small code example for each common question, with a link to a practice problem where one exists.
The questions usually arrive as a short code sample where you predict the output shape or value, or as a small task such as writing softmax for a batch of rows, followed by a question about an edge case such as overflow, ties or a shape that broadcasts by accident. If your loop also has a DataFrame round, our pandas interview questions guide covers it, and our ML coding interview questions guide lists the models you are most often asked to implement on top of NumPy. The code on this page assumes import numpy as np and runs on current NumPy.
What do interviewers ask about NumPy arrays and dtypes?
Interviewers ask why NumPy arrays are faster than Python lists and how the dtype of an array changes memory use and arithmetic. These questions check that you understand what NumPy does underneath, which helps you explain every later answer about speed.
Why is a NumPy array faster than a Python list?
A NumPy array is faster because it stores values of one data type in one contiguous block of memory, so an operation such as multiplication runs as a single loop in compiled code over raw numbers. A Python list stores pointers to separate Python objects, and a list comprehension runs the interpreter once per element and checks types at every step. The array also uses less memory, since each float64 value takes 8 bytes with no per-object overhead. The two lines below compute the same totals, and on arrays with millions of elements the second is many times faster. Practise building arrays with Create Arrays from Lists.
import numpy as np
prices = [3.0, 4.5, 2.0, 8.0]
qty = [2, 1, 4, 3]
totals_list = [p * q for p, q in zip(prices, qty)]
totals_array = np.array(prices) * np.array(qty)What is a dtype, and why does it matter?
A dtype is the data type that every element of an array shares, such as int64, float32 or bool, and it fixes how many bytes each value takes and how arithmetic behaves. NumPy infers it when you create an array, so a list of whole numbers gives an integer array and a single float in the list turns the whole array into float64. Fixed-size integers wrap around on overflow without any error, which is why the uint8 sum below prints 244 instead of 500, and assigning a float into an integer array silently truncates it. Deep learning code often uses float32 because it halves memory compared with float64. Zeros and Ones and Arange and Linspace practise creating arrays with a chosen shape.
ints = np.array([1, 2, 3])
floats = np.array([1, 2, 3.5])
pixels = np.array([250, 10], dtype=np.uint8)
print(ints.dtype, floats.dtype)
print(pixels + pixels)
print(pixels.astype(np.int32) + pixels)
print(np.zeros(1000, dtype=np.float32).nbytes)What shape and reshape questions do interviewers ask?
Interviewers ask how reshape, ravel and flatten differ and when NumPy returns a view instead of a copy. Interviewers ask these often because a shape or memory mistake in model code usually runs without an error and gives wrong numbers.
How do reshape, ravel and flatten differ?
The reshape method returns the same data with a new shape, and one dimension can be -1 so NumPy works it out from the total size. The ravel method returns a flat view of the array whenever the memory layout allows it and copies only when it must, while flatten always returns a new copy. That means writing into the result of ravel can change the original array, and writing into the result of flatten never does. Indexing with None, which is the same as np.newaxis, adds a dimension of size 1, which you need for broadcasting. Practise with Reshaping Arrays.
x = np.arange(12)
grid = x.reshape(3, 4)
pairs = x.reshape(-1, 2)
flat_view = grid.ravel()
flat_copy = grid.flatten()
flat_copy[0] = 100
print(grid[0, 0])
flat_view[0] = 100
print(grid[0, 0])
column = x[:, None]
print(pairs.shape, column.shape)When does NumPy return a view, and when a copy?
NumPy returns a view for basic slicing with integers and colons, and a copy for fancy indexing with an array of positions and for boolean masks. A view shares memory with the original array, so writing into it changes the original, which is a frequent source of bugs when a function modifies a slice it was given. The np.shares_memory function tells you which case you have, and calling copy gives you independent data when you need it. Most interviewers ask this with a short code sample and want you to predict the printed array. The Mutation Trap problem is built around it.
a = np.arange(6)
window = a[1:4]
window[0] = 99
print(a)
picked = a[[1, 2, 3]]
picked[0] = -1
print(a)
print(np.shares_memory(a, window), np.shares_memory(a, picked))
safe = a[1:4].copy()What indexing questions do interviewers ask in NumPy?
Interviewers ask how basic slicing, fancy indexing and boolean masks select data, and what each one returns. A good answer names the output shape and says whether the result is a view or a copy.
How do basic slicing, fancy indexing and boolean masks differ?
Basic slicing uses integers and slices and returns a view, fancy indexing uses arrays of integer positions and returns a copy in the order you asked for, and a boolean mask keeps the elements where a same-shaped boolean array is True and returns them as a flat copy. Passing two index arrays pairs them element by element, so m[[0, 2], [1, 3]] returns m[0, 1] and m[2, 3] instead of a 2 by 2 block, and np.ix_ is the way to get the block. Combine conditions with &, | and ~ and wrap each one in parentheses, because the Python keywords and and or raise a ValueError on arrays. The np.where function picks between two values per element, and assigning through a mask changes only the matching elements. Practise these with Basic Indexing, Fancy Indexing, Boolean Masking and Filter and Extract.
m = np.arange(12).reshape(3, 4)
row = m[1]
col = m[:, 2]
block = m[:2, 1:3]
pairs = m[[0, 2], [1, 3]]
sub_grid = m[np.ix_([0, 2], [1, 3])]
even_large = m[(m > 3) & (m % 2 == 0)]
labels = np.where(m > 5, "high", "low")
m[m > 9] = 0How does NumPy broadcasting work?
Broadcasting lets NumPy combine arrays of different shapes by lining them up from the right and stretching any dimension of size 1 to match the other array. Interviewers ask you to state the rules and then predict whether a pair of shapes works and what shape it produces.
What are the broadcasting rules?
NumPy compares two shapes from the last dimension backwards, and two dimensions are compatible when they are equal or when one of them is 1, in which case that array is stretched along the dimension without copying data. If one array has fewer dimensions, NumPy first pads its shape with 1s on the left. Shapes (3, 1) and (4,) therefore become (3, 1) and (1, 4) and broadcast to (3, 4). Shapes (3, 4) and (3,) raise a ValueError, because the last dimensions are 4 and 3, and the fix is to reshape the second array to (3, 1) with [:, None]. The same rule produces an outer sum or a table of pairwise differences from two vectors, which Outer Sum and Pairwise Differences practise.
col = np.arange(3).reshape(3, 1)
row = np.arange(4)
print((col + row).shape)
X = np.arange(12.0).reshape(3, 4)
offsets = np.array([10.0, 20.0, 30.0])
shifted = X + offsets[:, None]
outer = row[:, None] + row[None, :]
diffs = row[:, None] - row[None, :]
print(shifted.shape, outer.shape, diffs[0])How do you scale every row or every column of a matrix?
You scale columns by combining the matrix with a statistic computed over axis=0, and rows with a statistic computed over axis=1 and keepdims=True. Column statistics have shape (n_features,) and line up with an (n_samples, n_features) matrix from the right without any help. Row statistics have shape (n_samples,), which does not line up, so keepdims=True keeps them as (n_samples, 1) and lets them broadcast across each row. Practise both directions with Row Scaling and Column Scaling.
X = np.array([[1.0, 2.0, 3.0], [4.0, 6.0, 8.0]])
col_centered = X - X.mean(axis=0)
row_share = X / X.sum(axis=1, keepdims=True)
col_scaled = X / X.max(axis=0)
print(row_share.sum(axis=1))What does the axis argument mean in NumPy?
The axis argument names the dimension that an operation runs along, so sum(axis=0) on a 2D array adds down the rows and returns one value per column, and sum(axis=1) adds across each row and returns one value per row. For reductions such as sum, mean, max and argmax, the axis you pass disappears from the output shape unless you also pass keepdims=True, which leaves it with size 1. Functions that keep the shape, such as cumsum and np.sort, work along the axis instead of removing it. The value axis=-1 always means the last dimension, which is the usual choice for softmax over class scores. Practise reductions with Aggregation Functions and Row Extremes.
sales = np.array([[3, 5, 2], [4, 1, 6]])
print(sales.sum(axis=0))
print(sales.sum(axis=1))
print(sales.sum(axis=1, keepdims=True).shape)
print(sales.argmax(axis=1))
print(sales.cumsum(axis=1))How do you vectorise a Python loop in NumPy?
You vectorise a loop by rewriting the work done per element as operations on whole arrays, using arithmetic and comparisons for the body, np.where or np.maximum for branches, a reduction with an axis for a running sum, and broadcasting for nested loops over pairs. The np.vectorize function does not count as vectorisation, because it still calls your Python function once per element. In an interview it is fine to write the loop first, then replace it and check that both versions give the same output on a small input, as in the hinge loss below. The Vectorized Operations section of the NumPy Sheet practises this with problems such as Normalized Difference, Angle Features and Norm-Gated Linear Transform.
def hinge_loop(scores, labels):
total = 0.0
for s, y in zip(scores, labels):
margin = 1.0 - y * s
if margin > 0:
total += margin
return total / len(scores)
def hinge_vectorised(scores, labels):
return np.maximum(0.0, 1.0 - labels * scores).mean()
scores = np.array([0.8, -0.3, 2.0, 0.1])
labels = np.array([1.0, -1.0, 1.0, -1.0])
print(np.isclose(hinge_loop(scores, labels), hinge_vectorised(scores, labels)))What is the difference between @, np.dot and * in NumPy?
The * operator multiplies arrays element by element, while @ and np.dot compute matrix products that sum over the last axis of the first array and the second-to-last axis of the second. For 1D and 2D arrays, @ and np.dot give the same result. For arrays with three or more dimensions, @ treats the leading dimensions as a batch and multiplies matching matrices, while np.dot pairs every matrix in the first array with every matrix in the second, which is why the last line below prints (8, 2, 8, 4). The @ operator also rejects scalars, where np.dot accepts them, so @ is the safer default for matrix code. Practise with Matrix Multiply, and read the matrix multiplication lesson for the maths.
A = np.arange(6).reshape(2, 3)
B = np.arange(12).reshape(3, 4)
v = np.array([1, 2, 3])
print((A @ B).shape, np.array_equal(A @ B, np.dot(A, B)))
print((A * A).shape, A @ v, v @ v)
batch_a = np.ones((8, 2, 3))
batch_b = np.ones((8, 3, 4))
print((batch_a @ batch_b).shape)
print(np.dot(batch_a, batch_b).shape)How does np.einsum work?
The np.einsum function describes an operation with one letter per axis, sums over every letter that appears in the inputs but not in the output, and orders the output axes by the letters after the arrow. The string ij,jk->ik is a matrix product, ii-> is a trace, ij->ji is a transpose, and bqd,bkd->bqk computes attention scores for a batch of queries and keys without any manual transpose. Interviewers ask about it because one string replaces a chain of transposes and reshapes, and you should be able to translate any einsum back into @ and sum to check it. Named-Dimension Batched Attention Scores practises the same index notation, and Scaled Dot-Product Attention uses the scores in a full attention step.
rng = np.random.default_rng(0)
A = rng.normal(size=(2, 3))
B = rng.normal(size=(3, 4))
S = rng.normal(size=(3, 3))
q = rng.normal(size=(2, 5, 8))
k = rng.normal(size=(2, 7, 8))
print(np.allclose(np.einsum("ij,jk->ik", A, B), A @ B))
print(np.isclose(np.einsum("ii->", S), np.trace(S)))
scores = np.einsum("bqd,bkd->bqk", q, k)
print(scores.shape, np.allclose(scores, q @ k.transpose(0, 2, 1)))How should you generate random numbers in NumPy?
You should create a Generator with np.random.default_rng and a seed, and draw every random number from it, which is the approach the NumPy documentation recommends over the older global functions such as np.random.seed and np.random.rand. A Generator keeps its own state, so two parts of a program cannot change each other's sequence, and the same seed gives the same numbers on every run. Pass the generator into functions that need randomness instead of seeding inside them, so a test can fix the seed from outside. A shuffled train and test split is a permutation of row indices, and dropout masks and mini-batches come from random and choice. Practise with Random Array Generation.
rng = np.random.default_rng(42)
weights = rng.normal(0.0, 0.02, size=(4, 3))
order = rng.permutation(10)
batch = rng.choice(100, size=5, replace=False)
keep = rng.random(10) < 0.8
same = np.random.default_rng(42).normal(0.0, 0.02, size=(4, 3))
print(np.array_equal(weights, same))How do you keep NumPy code numerically stable?
You keep NumPy code numerically stable by avoiding operations that overflow or lose precision, such as calling np.exp on large scores or taking the log of a number that has rounded to zero. Interviewers ask about softmax and logsumexp most often, because both appear in almost every classification model and both break on large inputs when written naively.
How do you write a numerically stable softmax?
You write a numerically stable softmax by subtracting the maximum of each row before calling np.exp, which leaves the result unchanged because the shift cancels between the numerator and the denominator. Without the shift, np.exp(1000.0) overflows to inf and the division returns nan. Passing axis=-1 with keepdims=True lets the same function handle a single vector and a batch of rows, and the two rows below give identical probabilities because they differ only by a constant. Practise it with the numerically stable softmax problem.
def softmax(z, axis=-1):
shifted = z - z.max(axis=axis, keepdims=True)
e = np.exp(shifted)
return e / e.sum(axis=axis, keepdims=True)
z = np.array([[1000.0, 1001.0, 1002.0], [1.0, 2.0, 3.0]])
print(softmax(z).round(3))What is logsumexp, and when do you need it?
The logsumexp function computes the log of a sum of exponentials as the maximum plus the log of the sum of exponentials after subtracting that maximum, which stays finite for large scores. You need it whenever you work with log-probabilities, since log-softmax equals the scores minus their logsumexp and cross-entropy from logits uses it to avoid taking the log of a probability that rounded to zero. Other habits interviewers like to hear are np.log1p for the log of 1 plus a small number, clipping probabilities away from 0 before a log, and comparing floats with np.isclose instead of ==. Practise the loss with Cross-Entropy Loss, and read the cross-entropy lesson for the maths.
def logsumexp(z, axis=-1):
m = z.max(axis=axis, keepdims=True)
return (m + np.log(np.exp(z - m).sum(axis=axis, keepdims=True))).squeeze(axis)
def cross_entropy(logits, labels):
log_probs = logits - logsumexp(logits)[:, None]
return -log_probs[np.arange(len(labels)), labels].mean()
logits = np.array([[1000.0, 1001.0, 1002.0], [2.0, 0.5, -1.0]])
labels = np.array([2, 0])
print(logsumexp(logits))
print(cross_entropy(logits, labels))
print(0.1 + 0.2 == 0.3, np.isclose(0.1 + 0.2, 0.3))How do you write common ML operations in NumPy?
You write common ML operations in NumPy by combining indexing, broadcasting and reductions, and interviewers most often ask for one-hot encoding, pairwise distances, feature normalisation and top-k selection. Each of these has a loop version and a vectorised version, and the interviewer usually wants the vectorised one.
How do you one-hot encode labels in NumPy?
You one-hot encode integer labels by indexing the rows of an identity matrix with the labels, or by creating an array of zeros and setting one entry per row with fancy indexing. Both give an array of shape (n_samples, n_classes). The second version avoids building a full identity matrix, which matters when the number of classes is large. Practise it with the one-hot encoding problem.
labels = np.array([2, 0, 1, 2])
n_classes = 3
one_hot = np.eye(n_classes)[labels]
manual = np.zeros((labels.size, n_classes))
manual[np.arange(labels.size), labels] = 1.0
print(np.array_equal(one_hot, manual))How do you compute pairwise distances without a loop?
You compute every pairwise Euclidean distance between the rows of X and Y either by broadcasting their difference to shape (n, m, d) or by expanding the squared distance into the squared norm of x plus the squared norm of y minus twice their dot product. The expanded form replaces the large three-dimensional array with one matrix product, so it is the usual answer for large inputs. Rounding can make a tiny squared distance slightly negative in that form, so clip it at zero before the square root. Practise with Euclidean Distance, and use the result in KNN Classifier.
def pairwise_broadcast(X, Y):
diff = X[:, None, :] - Y[None, :, :]
return np.sqrt((diff ** 2).sum(axis=-1))
def pairwise_matmul(X, Y):
sq = (X ** 2).sum(axis=1)[:, None] + (Y ** 2).sum(axis=1)[None, :] - 2 * X @ Y.T
return np.sqrt(np.maximum(sq, 0.0))
rng = np.random.default_rng(1)
X = rng.normal(size=(5, 3))
Y = rng.normal(size=(4, 3))
print(np.allclose(pairwise_broadcast(X, Y), pairwise_matmul(X, Y)))How do you normalise features in NumPy?
You standardise features by subtracting each column's mean and dividing by its standard deviation over axis=0, you min-max scale with each column's minimum and range, and you scale rows to unit length by dividing by np.linalg.norm with axis=1 and keepdims=True. Add a small epsilon to each denominator, because a constant column has a standard deviation and a range of zero. In a real pipeline the statistics come from the training set and are reused on the test set, and interviewers often ask about that. Practise with Normalize Columns, use unit rows in Cosine Similarity, and clip outliers with Winsorize.
X = np.array([[1.0, 200.0, 5.0], [2.0, 300.0, 5.0], [3.0, 400.0, 5.0]])
eps = 1e-8
standardised = (X - X.mean(axis=0)) / (X.std(axis=0) + eps)
min_max = (X - X.min(axis=0)) / (X.max(axis=0) - X.min(axis=0) + eps)
unit_rows = X / np.linalg.norm(X, axis=1, keepdims=True)
print(standardised[:, 2], np.linalg.norm(unit_rows, axis=1))How do you find the top k values with argsort?
You find the top k values by sorting indices with np.argsort on the negated scores and taking the first k, or with np.argpartition when the array is large, since it moves the k largest values to the end in linear time and you then sort only those k. The argsort function sorts in ascending order, so negating the scores or reversing the result gives descending order. Pass kind="stable" when ties should keep their original order. For a batch of rows, pass axis=1 and gather the values with np.take_along_axis. Practise with Sort and Argsort.
scores = np.array([0.2, 0.9, 0.4, 0.7, 0.1, 0.8])
k = 3
top_sorted = np.argsort(-scores)[:k]
part = np.argpartition(scores, -k)[-k:]
top_part = part[np.argsort(-scores[part])]
print(top_sorted, top_part)
batch = np.array([[0.1, 0.5, 0.3], [0.9, 0.2, 0.4]])
top2 = np.argsort(-batch, axis=1)[:, :2]
print(np.take_along_axis(batch, top2, axis=1))Where can you find NumPy practice problems?
You can find NumPy practice problems in the TensorTonic NumPy Sheet, which has 25 problems in five sections that run from array creation through indexing and slicing, broadcasting and vectorized operations to array utilities such as winsorizing, tiling, differencing, concatenation and correlation. Every solution runs against tests, so you see straight away whether your array matches the expected output.
- Work through the NumPy Sheet in order, because the broadcasting and vectorized operations sections build on the indexing section before them.
- Rewrite every snippet on this page from memory, write down the shape you expect after each line, and then run it to check.
- Continue with the Linear Algebra study plan, which has 30 problems in five sections, all written in NumPy, and read our guide to learning linear algebra for machine learning alongside it.
- Move on to Cracking ML, which builds classic models from scratch in NumPy, so you use these operations inside full algorithms.
- Take the 60-minute Pandas + NumPy Sprint to practise under a timer, and then the free Data Science Mock Interview to see NumPy next to the other rounds.
What mistakes do candidates make with NumPy in interviews?
The mistakes interviewers see most often are shapes that broadcast by accident, changes to a view that leak into the original array, and integer dtypes that silently truncate or overflow. Each of them produces output without an error, which is why interviewers ask about them directly.
- Subtracting a (3, 1) column of predictions from a (3,) vector of targets broadcasts to a (3, 3) matrix, so a mean squared error computed that way averages nine numbers instead of three, as the example below shows.
- Changing a slice inside a function changes the caller's array, because basic slicing returns a view, so call copy when the function should not modify its input.
- Assigning 0.9 into an integer array stores 0, and adding two uint8 arrays wraps around above 255, so check the dtype with
arr.dtypebefore you write results into an existing array. - Forgetting keepdims=True on a row reduction gives a shape of (n,) that fails to broadcast against an (n, d) matrix, or broadcasts against the wrong axis when n equals d.
- Using the Python keywords and and or between boolean arrays raises a ValueError, and the fix is & and | with parentheses around each condition.
- Looping over rows with a for loop when a single expression with an axis argument would do the same work is the slowest way to touch every row, and interviewers usually ask you to remove it.
- Comparing floats with == fails on values such as 0.1 + 0.2, so use np.isclose or np.allclose in checks and tests.
y = np.array([1.0, 2.0, 3.0])
pred = np.array([[1.5], [2.0], [2.5]])
print((y - pred).shape)
wrong = ((y - pred) ** 2).mean()
right = ((y - pred.ravel()) ** 2).mean()Practice this
- Boolean Masking drills selecting data with boolean conditions, which most data cleaning questions need.
- Normalize Columns combines axis arguments and broadcasting in the way feature scaling questions do.
- Sort and Argsort drills the index sorting behind top-k, ranking and nearest neighbour questions.
Then take the 60-minute Pandas + NumPy Sprint to answer these questions under a timer.
Do machine learning interviews ask NumPy questions?
Many machine learning and data scientist interviews include NumPy, either as direct questions about arrays, shapes and broadcasting or inside an ML coding round where you implement something such as softmax, k-means or linear regression on NumPy arrays. Interviewers use it to see whether you think in whole-array operations or fall back on Python loops. Ask your recruiter whether the coding round allows NumPy and whether you can look up documentation during it.
Which NumPy functions should you know for an interview?
You should know array, zeros, arange, linspace and reshape for creating arrays, slicing, fancy indexing, boolean masks and where for selecting data, sum, mean, std, max and argmax with an axis for reductions, and matmul, einsum and linalg.norm for linear algebra. Add argsort, argpartition, clip, concatenate, stack and the random Generator methods normal, choice and permutation. A typical interview task combines several of these, so practise chaining them on small arrays until you can predict every output shape.
How much NumPy do you need before learning PyTorch?
You need enough NumPy to create and reshape arrays, index and mask them, predict the result of broadcasting and reduce along an axis, because PyTorch tensors follow the same rules. Most NumPy code you write for an interview carries over to PyTorch once you swap np for torch and axis for dim. Autograd and the GPU are the parts PyTorch adds on top.
Should you prepare NumPy or pandas for a data science interview?
You should prepare both if the role involves Python, because pandas questions cover tables, joins and grouping while NumPy questions cover arrays and the maths inside models. Vectorised pandas code also relies on NumPy functions such as numpy.where and numpy.select. Analytics roles lean towards pandas, and modeling roles ask more NumPy, so ask the recruiter which one the coding round uses.
Is there a timed NumPy practice test?
Yes, TensorTonic has a 60-minute Pandas + NumPy Sprint in its Data Scientist track that covers pandas and NumPy fundamentals. The free Data Science Mock Interview is a 75-minute sitting with five problems that mix SQL, pandas, NumPy, math and ML, and it ends with a scored report.