ML coding interviews

Setup#

pip install numpy

Everything below runs on a plain CPU in under two seconds, with no downloads. The three algorithms here — k-means, attention, backprop — are the three most commonly asked, and each one carries a trap that interviewers deliberately probe.

For the DSA flavour of this round, this site's /leetcode/ archive has 1,055 worked solutions, filterable by topic. The rest of this lesson is the from-scratch flavour.

1. K-means, with the trap included#

The task as asked: "Implement k-means. You may use NumPy." Expected in roughly 20 minutes.

K-means repeats two steps: assign every point to its nearest centroid, then move each centroid to the mean of its points. The trap is what nobody mentions: where the centroids start.

import numpy as np

rng = np.random.default_rng(0)

# 30 points in 2D, drawn around three separated centres
data = np.vstack([
    rng.normal([0, 0], 0.5, (10, 2)),
    rng.normal([5, 5], 0.5, (10, 2)),
    rng.normal([0, 5], 0.5, (10, 2)),
])


def kmeans(data, k, seed, steps=20):
    rng = np.random.default_rng(seed)
    centroids = data[rng.choice(len(data), k, replace=False)]  # start on real points
    for _ in range(steps):
        # distance from every point to every centroid, no Python loops
        d = np.linalg.norm(data[:, None, :] - centroids[None, :, :], axis=2)
        labels = d.argmin(axis=1)                  # assign: nearest centroid wins
        new = np.array([data[labels == j].mean(axis=0) for j in range(k)])
        if np.allclose(new, centroids):            # centroids stopped moving
            break
        centroids = new
    inertia = (d.min(axis=1) ** 2).sum()           # the number k-means shrinks
    return centroids, inertia


results = [kmeans(data, k=3, seed=s) for s in range(5)]
for s, (_, inertia) in enumerate(results):
    print(f"seed {s}  final inertia {inertia:8.3f}")

best, _ = min(results, key=lambda r: r[1])
print("best centroids found:")
print(np.round(best[np.argsort(best[:, 0])], 2))   # sorted so the order is stable

seed 0  final inertia  148.568
seed 1  final inertia  148.568
seed 2  final inertia  124.485
seed 3  final inertia  124.415
seed 4  final inertia   10.772
best centroids found:
[[-0.17 -0.01]
 [ 0.1   5.26]
 [ 4.9   5.16]]

Read that output slowly, because it is the whole interview. Five random starts on the same easy data gave three different answers, and four of the five were bad. Only seed 4 found the true clusters, with inertia dropping from 148 to 10.8.

Inertia is the sum of squared distances from each point to its assigned centroid — the quantity k-means minimises. Each run only ever lowers it, but runs get stuck in local minima: arrangements no single step can improve, that are still wrong. Two centroids land in one blob and split it, while one centroid stretches across two blobs.

This is the follow-up question, pre-answered: "How do you make k-means reliable?" Run it several times and keep the lowest inertia, or start smarter. Smarter starting is k-means++, which picks initial centroids spread far apart — it is what libraries do by default. The concepts live in the clustering lesson.

The other line worth narrating is the distance computation. data[:, None, :] - centroids[None, :, :] uses broadcasting — NumPy stretching shapes (30, 1, 2) and (1, 3, 2) into (30, 3, 2) — to get every point-to-centroid distance without a Python loop. Say the shapes out loud as you write them. Interviewers grade that.

2. Scaled dot-product attention#

The task as asked: "Write the attention function: given Q, K, V, return the output. Then add a causal mask." Expected in roughly 15 minutes.

Attention lets each token in a sequence build its output as a weighted mix of every token's information, with the weights computed from relevance. If the idea itself is new to you, read attention first — this is the implementation.

import numpy as np

rng = np.random.default_rng(0)

seq_len, d_k = 4, 8                        # four tokens, eight dimensions each
Q = rng.normal(size=(seq_len, d_k))        # what each token is looking for
K = rng.normal(size=(seq_len, d_k))        # what each token offers to be found by
V = rng.normal(size=(seq_len, d_k))        # what each token hands over if chosen


def softmax(x):
    x = x - x.max(axis=-1, keepdims=True)  # subtract the max: overflow-proof, same answer
    e = np.exp(x)
    return e / e.sum(axis=-1, keepdims=True)


scores = Q @ K.T / np.sqrt(d_k)            # every token scored against every token
weights = softmax(scores)                  # each row becomes a probability distribution
out = weights @ V                          # each output row is a weighted mix of V rows

print("attention weights, one row per token:")
print(np.round(weights, 3))
print("row sums:", np.round(weights.sum(axis=1), 3))
print("output shape:", out.shape)

# the causal mask a decoder adds: a token may not look at tokens after itself
mask = np.triu(np.ones((seq_len, seq_len)), k=1).astype(bool)
masked = softmax(np.where(mask, -1e9, scores))
print("with a causal mask (upper triangle forced to zero):")
print(np.round(masked, 3))

attention weights, one row per token:
[[0.302 0.472 0.081 0.145]
 [0.329 0.069 0.247 0.354]
 [0.32  0.388 0.233 0.059]
 [0.102 0.031 0.684 0.183]]
row sums: [1. 1. 1. 1.]
output shape: (4, 8)
with a causal mask (upper triangle forced to zero):
[[1.    0.    0.    0.   ]
 [0.827 0.173 0.    0.   ]
 [0.34  0.413 0.247 0.   ]
 [0.102 0.031 0.684 0.183]]

Three graded details hide in eleven lines of core code.

The max subtraction in softmax. np.exp(1000) overflows to infinity and the weights become nan. Subtracting the row maximum changes nothing mathematically — softmax only cares about differences — and makes overflow impossible. Interviewers feed large scores on purpose to see if your softmax survives.

The scaling by the square root of the dimension. With larger d_k, dot products grow, softmax saturates towards one-hot, and gradients die. Dividing by the square root of d_k keeps score variance near 1. Saying that sentence is the difference between "memorised the formula" and "understands it".

The mask uses a large negative number, not zero. Zeroing scores does not remove a token: a score of 0 still earns weight after softmax. Setting masked scores to a huge negative number drives their weight to zero. Check the first row of the masked output — the first token can only attend to itself, so its row reads 1.0 and zeros.

3. Backprop on a tiny network#

The task as asked: "Train a small network on XOR without any framework. Write the gradients by hand." Expected in roughly 25 minutes — this is the hardest of the three.

XOR is the classic because a single layer cannot solve it — no straight line separates the classes — so your hidden layer and its gradients must genuinely work. If the chain-rule story is fuzzy, backpropagation teaches it fully; this is the interview version.

import numpy as np

rng = np.random.default_rng(42)

X = np.array([[0, 0], [0, 1], [1, 0], [1, 1]], dtype=float)
y = np.array([[0], [1], [1], [0]], dtype=float)    # XOR: no straight line separates it

W1 = rng.normal(0, 1, (2, 4)); b1 = np.zeros(4)
W2 = rng.normal(0, 1, (4, 1)); b2 = np.zeros(1)
lr = 1.0


def sigmoid(z):
    return 1 / (1 + np.exp(-z))


for step in range(3001):
    # forward pass, keeping every intermediate - backprop will need them
    h = np.tanh(X @ W1 + b1)
    p = sigmoid(h @ W2 + b2)
    loss = -(y * np.log(p) + (1 - y) * np.log(1 - p)).mean()

    # backward pass: chain rule, one layer at a time, top to bottom
    dz2 = (p - y) / len(X)          # sigmoid + cross-entropy collapses to this
    dW2 = h.T @ dz2; db2 = dz2.sum(0)
    dh = dz2 @ W2.T                 # send the gradient through W2 to the hidden layer
    dz1 = dh * (1 - h ** 2)         # tanh'(z) = 1 - tanh(z)^2, and h is tanh(z)
    dW1 = X.T @ dz1; db1 = dz1.sum(0)

    for param, grad in ((W1, dW1), (b1, db1), (W2, dW2), (b2, db2)):
        param -= lr * grad
    if step % 1000 == 0:
        print(f"step {step:4d}  loss {loss:.4f}")

print("predictions for [00, 01, 10, 11]:", np.round(p.ravel(), 3))

step    0  loss 0.8299
step 1000  loss 0.0024
step 2000  loss 0.0011
step 3000  loss 0.0007
predictions for [00, 01, 10, 11]: [0.    0.999 0.999 0.001]

The loss falls three orders of magnitude and the four predictions land on XOR's truth table. Two lines carry almost all the marks.

dz2 = (p - y) / len(X) — the gradient of cross-entropy loss through a sigmoid collapses to prediction minus target. Deriving or at least stating this collapse is a standard follow-up. The division by len(X) matches the .mean() in the loss; forget it and your gradient is four times too large.

dz1 = dh * (1 - h ** 2) — the chain rule passing through tanh. The elegance is that h was already computed in the forward pass, so the derivative costs almost nothing. That reuse of forward intermediates is backpropagation.

If your interview allows a shape check, one sanity rule catches most bugs: every gradient has the same shape as its parameter. dW1.shape == W1.shape, always. Check it before running.

Common mistakes across all three#

Starting to type immediately. State the algorithm in two sentences first. It costs 20 seconds and buys you a plan, plus visible marks.

No tiny test. Run k-means on ten points, attention on four tokens, the network on XOR. A candidate who invents a three-line test case reads as someone who has debugged real models.

Looping where broadcasting works. Triple-nested Python loops for distances are correct and score poorly. Knowing the vectorised form is part of what the round measures.

Dividing by zero on an empty cluster. In k-means, a centroid can end up with no points, and mean of nothing is nan. Mentioning the case — "I would reseed an empty centroid to a random point" — earns more than silently hoping.

Forgetting numerical stability. The unstable softmax is the single most common failure in the attention question. Write the max subtraction by reflex.

Try it yourself#

Rebuild each of the three from a blank file, without looking, one per day for three days. Then do the standard variations: k-means with the empty-cluster fix, attention with two heads (split d_k in half, run twice, concatenate), and the XOR network with ReLU hidden units instead of tanh — which changes exactly one forward line and one backward line.

What to learn next#

• ML theory questions, with answers — the spoken round that usually follows this one.
• The LeetCode archive — 1,055 worked solutions for the DSA flavour of the coding round.
• Backpropagation — the full lesson behind the hardest of the three tasks.