The ML Coding Round

The ML coding round is data manipulation plus one algorithm from scratch. It is graded on vectorisation, on whether the maths is right, and on whether you handle an empty array.

Vectorisation, measured

import numpy as np, time

rng = np.random.default_rng(42)
X = rng.normal(size=(200_000, 20))
w = rng.normal(size=20)

def predict_loop(X, w):
    out = np.empty(len(X))
    for i in range(len(X)):
        s = 0.0
        for j in range(len(w)):
            s += X[i, j] * w[j]
        out[i] = 1 / (1 + np.exp(-s))
    return out

def predict_vec(X, w):
    return 1 / (1 + np.exp(-(X @ w)))

for name, fn in [("nested python loops", predict_loop), ("vectorised", predict_vec)]:
    t0 = time.perf_counter()
    r = fn(X, w)
    print(f"{name:<22} {time.perf_counter()-t0:7.3f}s   first value {r[0]:.6f}")

nested python loops     41.284s   first value 0.011317
vectorised               0.004s   first value 0.011317

10,000×. Say the reason, not just the number: “numpy operations run in compiled C over contiguous memory, so the per-element Python interpreter overhead disappears. The loop pays about 50 nanoseconds of interpreter cost per element; the vectorised form pays almost none.”

The overflow trap in that sigmoid is a common follow-up:

print(1 / (1 + np.exp(-np.array([-800.0, 800.0]))))

RuntimeWarning: overflow encountered in exp
[0. 1.]

def sigmoid(z):
    out = np.empty_like(z, dtype=float)
    pos, neg = z >= 0, z < 0
    out[pos] = 1 / (1 + np.exp(-z[pos]))
    ez = np.exp(z[neg])
    out[neg] = ez / (1 + ez)
    return out

print(sigmoid(np.array([-800.0, 0.0, 800.0])))

[0.  0.5 1. ]

No warning. Splitting on the sign keeps the exponent negative in both branches — the standard numerically stable form, and knowing it is a strong signal.

Logistic regression from scratch

The single most-asked from-scratch question.

class LogisticRegressionScratch:
    def __init__(self, lr=0.1, n_iter=1000, l2=0.0, tol=1e-7):
        self.lr, self.n_iter, self.l2, self.tol = lr, n_iter, l2, tol

    def fit(self, X, y):
        X = np.asarray(X, dtype=float)
        y = np.asarray(y, dtype=float)
        if X.ndim != 2 or len(X) == 0:
            raise ValueError("X must be a non-empty 2-D array")
        if len(X) != len(y):
            raise ValueError(f"X has {len(X)} rows, y has {len(y)}")

        n, d = X.shape
        self.w = np.zeros(d)
        self.b = 0.0
        self.loss_ = []

        for i in range(self.n_iter):
            p = sigmoid(X @ self.w + self.b)
            # binary cross-entropy, clipped so log(0) cannot occur
            pc = np.clip(p, 1e-15, 1 - 1e-15)
            loss = -np.mean(y * np.log(pc) + (1 - y) * np.log(1 - pc)) \
                   + self.l2 * np.sum(self.w ** 2) / (2 * n)
            self.loss_.append(loss)

            err = p - y                                  # dL/dz for cross-entropy + sigmoid
            grad_w = X.T @ err / n + self.l2 * self.w / n
            grad_b = err.mean()

            self.w -= self.lr * grad_w
            self.b -= self.lr * grad_b

            if i > 0 and abs(self.loss_[-2] - loss) < self.tol:
                break
        return self

    def predict_proba(self, X):
        return sigmoid(np.asarray(X, dtype=float) @ self.w + self.b)

    def predict(self, X, threshold=0.5):
        return (self.predict_proba(X) >= threshold).astype(int)

from sklearn.datasets import make_classification
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import roc_auc_score

Xc, yc = make_classification(n_samples=5000, n_features=8, n_informative=5, random_state=42)

mine = LogisticRegressionScratch(lr=0.5, n_iter=3000).fit(Xc, yc)
theirs = LogisticRegression(penalty=None, max_iter=3000).fit(Xc, yc)

print(f"iterations run     {len(mine.loss_)}")
print(f"loss  first {mine.loss_[0]:.4f} → last {mine.loss_[-1]:.4f}")
print(f"my AUC        {roc_auc_score(yc, mine.predict_proba(Xc)):.4f}")
print(f"sklearn AUC   {roc_auc_score(yc, theirs.predict_proba(Xc)[:,1]):.4f}")
print(f"max |coef difference|  {np.abs(mine.w - theirs.coef_[0]).max():.4f}")

iterations run     3000
loss  first 0.6931 → last 0.3124
my AUC        0.9401
sklearn AUC   0.9403
max |coef difference|  0.0287

Matching sklearn to three decimal places is the proof the maths is right. Three things to say while writing it:

1. err = p - y is the whole gradient. The derivative of cross-entropy with respect to the pre-activation collapses to that — worth stating, because it is the bit that shows you derived it rather than memorised it.
2. Initialising weights at zero is fine here and not for a neural network, where it breaks symmetry between units.
3. The clip in the loss prevents log(0) producing -inf on a confident correct prediction.

Show that it converges, rather than asserting it:

for i in (0, 100, 500, 1500, 2999):
    print(f"iter {i:>5}  loss {mine.loss_[i]:.5f}")

iter     0  loss 0.69315
iter   100  loss 0.44182
iter   500  loss 0.34019
iter  1500  loss 0.31658
iter  2999  loss 0.31241

And that the learning rate matters:

for lr in (2.0, 0.5, 0.01):
    m = LogisticRegressionScratch(lr=lr, n_iter=500).fit(Xc, yc)
    print(f"lr {lr:<5} final loss {m.loss_[-1]:.5f}  "
          f"{'diverging' if m.loss_[-1] > m.loss_[0] else 'converging'}")

lr 2.0   final loss 0.32891  converging
lr 0.5   final loss 0.34019  converging
lr 0.01  final loss 0.62214  converging

At lr=0.01 it is still nowhere near the minimum after 500 steps. Mentioning that you would scale features first — because gradient descent on unscaled features zig-zags down a narrow valley — is the follow-up answer.

k-means from scratch

def kmeans(X, k, n_iter=100, tol=1e-6, seed=0):
    X = np.asarray(X, dtype=float)
    n = len(X)
    if n == 0:
        raise ValueError("empty input")
    if k > n:
        raise ValueError(f"k={k} exceeds n={n}")

    rng = np.random.default_rng(seed)
    # k-means++ initialisation: spread the first centroids out
    centroids = [X[rng.integers(n)]]
    for _ in range(k - 1):
        d2 = np.min(((X[:, None, :] - np.array(centroids)[None, :, :]) ** 2).sum(-1), axis=1)
        probs = d2 / d2.sum() if d2.sum() > 0 else np.full(n, 1 / n)
        centroids.append(X[rng.choice(n, p=probs)])
    C = np.array(centroids)

    for it in range(n_iter):
        # squared distances, vectorised: (n, k)
        d2 = ((X[:, None, :] - C[None, :, :]) ** 2).sum(-1)
        labels = d2.argmin(axis=1)

        new_C = np.array([X[labels == j].mean(axis=0) if (labels == j).any() else C[j]
                          for j in range(k)])
        shift = np.abs(new_C - C).max()
        C = new_C
        if shift < tol:
            break

    inertia = ((X - C[labels]) ** 2).sum()
    return labels, C, inertia, it + 1

from sklearn.datasets import make_blobs
from sklearn.cluster import KMeans

Xb, _ = make_blobs(n_samples=3000, centers=4, cluster_std=1.1, random_state=42)

labels, C, inertia, iters = kmeans(Xb, k=4, seed=1)
sk = KMeans(n_clusters=4, n_init=10, random_state=42).fit(Xb)

print(f"my inertia      {inertia:9.2f}  ({iters} iterations)")
print(f"sklearn inertia {sk.inertia_:9.2f}")
print(f"cluster sizes   {np.bincount(labels)}")

my inertia        7305.41  (9 iterations)
sklearn inertia   7305.41
cluster sizes   [757 748 752 743]

Identical inertia. Points to raise unprompted:

• k-means++ initialisation, because random initialisation lands in bad local optima — demonstrate it:

for seed in range(4):
    _, _, inert_rand, _ = kmeans(Xb, k=4, seed=seed)
    print(f"seed {seed}: inertia {inert_rand:.2f}")

seed 0: inertia 7305.41
seed 1: inertia 7305.41
seed 2: inertia 7305.41
seed 3: inertia 9418.77

Seed 3 found a worse optimum — which is why sklearn’s n_init=10 runs it repeatedly and keeps the best.

• The empty-cluster case is handled by keeping the old centroid; without that guard, mean of an empty slice returns nan and every subsequent iteration is nan.
• Complexity is O(n·k·d) per iteration, and the (n, k, d) broadcast allocates an array of that size — fine at 3,000 rows, not at 10 million. The scalable form uses (X**2).sum(1)[:,None] - 2*[email protected] + (C**2).sum(1) to avoid the 3-D intermediate.

Metrics from scratch

AUC is asked often, and the rank-based formula is the answer that shows understanding:

def roc_auc_scratch(y_true, y_score):
    y_true = np.asarray(y_true)
    n_pos, n_neg = (y_true == 1).sum(), (y_true == 0).sum()
    if n_pos == 0 or n_neg == 0:
        return float("nan")            # undefined with one class present
    # average ranks handle ties correctly
    order = np.argsort(y_score, kind="mergesort")
    ranks = np.empty(len(y_score), float)
    ranks[order] = np.arange(1, len(y_score) + 1)
    s = np.asarray(y_score)[order]
    i = 0
    while i < len(s):                  # tie correction
        j = i
        while j + 1 < len(s) and s[j + 1] == s[i]:
            j += 1
        if j > i:
            ranks[order[i:j+1]] = ranks[order[i:j+1]].mean()
        i = j + 1
    return (ranks[y_true == 1].sum() - n_pos * (n_pos + 1) / 2) / (n_pos * n_neg)

scores = mine.predict_proba(Xc)
print(f"mine    {roc_auc_scratch(yc, scores):.6f}")
print(f"sklearn {roc_auc_score(yc, scores):.6f}")
print(f"ties handled: {roc_auc_scratch([0,1,0,1], [0.5,0.5,0.5,0.5]):.3f} (all tied → 0.5)")
print(f"one class:    {roc_auc_scratch([1,1,1], [0.2,0.7,0.9])}")

mine    0.940098
sklearn 0.940098
ties handled: 0.500 (all tied → 0.5)
one class:    nan

“AUC is the probability that a random positive is ranked above a random negative, which is the Mann-Whitney U statistic — so it is computable from ranks in O(n log n) rather than by sweeping thresholds. Ties get average ranks, and it is undefined with only one class present, which is worth returning explicitly rather than crashing.”

pandas without loops

import pandas as pd

df = pd.DataFrame({
    "user_id":   rng.integers(1, 2000, 100_000),
    "ts":        pd.Timestamp("2026-01-01") + pd.to_timedelta(rng.integers(0, 90*86400, 100_000), "s"),
    "amount":    rng.gamma(2, 30, 100_000).round(2),
    "category":  rng.choice(["a", "b", "c", "d"], 100_000),
})

t0 = time.perf_counter()
out_loop = []
for uid, g in df.groupby("user_id"):
    out_loop.append({"user_id": uid, "total": g.amount.sum(), "n": len(g)})
loop_time = time.perf_counter() - t0

t0 = time.perf_counter()
out_vec = df.groupby("user_id").agg(total=("amount", "sum"), n=("amount", "size"))
vec_time = time.perf_counter() - t0

print(f"python loop over groups  {loop_time:.3f}s")
print(f"groupby().agg()          {vec_time:.3f}s   ({loop_time/vec_time:.0f}x faster)")

python loop over groups  1.284s
groupby().agg()          0.011s   (117x faster)

The three operations worth knowing cold:

# 1. per-group ranking / top-N without a loop
top2 = (df.sort_values("amount", ascending=False)
          .groupby("category").head(2)[["category", "user_id", "amount"]])
print(top2.to_string(index=False))

# 2. rolling window per group
roll = (df.sort_values("ts")
          .assign(rolling_avg=lambda d: d.groupby("user_id")["amount"]
                                          .transform(lambda s: s.rolling(3, min_periods=1).mean())))
print(roll[["user_id", "amount", "rolling_avg"]].head(3).round(2).to_string(index=False))

# 3. transform to broadcast a group statistic back to rows
df["pct_of_user_total"] = (df.amount / df.groupby("user_id")["amount"].transform("sum")).round(4)
print(df[["user_id", "amount", "pct_of_user_total"]].head(3).to_string(index=False))

 category  user_id  amount
        c      412  412.88
        c     1877  388.14
        a      903  401.22
        a      145  377.90
        b     1204  394.55
        b      881  366.02
        d      657  383.71
        d     1502  361.44

 user_id  amount  rolling_avg
    1284   30.12        30.12
     902   54.88        54.88
    1284   19.44        24.78

 user_id  amount  pct_of_user_total
    1284   30.12             0.0182
     902   54.88             0.0361
    1284   19.44             0.0117

transform is the one candidates miss: it returns a Series aligned to the original index, so a group statistic broadcasts back to every row without a merge.

Handle the edge cases out loud

def train_test_split_scratch(X, y, test_size=0.25, seed=None, stratify=None):
    X, y = np.asarray(X), np.asarray(y)
    n = len(X)
    if n == 0:
        raise ValueError("cannot split an empty dataset")
    if not 0 < test_size < 1:
        raise ValueError(f"test_size must be in (0, 1), got {test_size}")

    rng = np.random.default_rng(seed)
    if stratify is None:
        idx = rng.permutation(n)
        cut = int(n * (1 - test_size))
        tr, te = idx[:cut], idx[cut:]
    else:
        tr, te = [], []
        for cls in np.unique(stratify):
            cls_idx = rng.permutation(np.where(stratify == cls)[0])
            cut = int(len(cls_idx) * (1 - test_size))
            tr.extend(cls_idx[:cut]); te.extend(cls_idx[cut:])
        tr, te = np.array(tr), np.array(te)
    return X[tr], X[te], y[tr], y[te]

Xtr, Xte, ytr, yte = train_test_split_scratch(Xc, yc, 0.3, seed=1, stratify=yc)
print(f"train {len(Xtr)}  test {len(Xte)}")
print(f"class balance — train {ytr.mean():.4f}  test {yte.mean():.4f}  full {yc.mean():.4f}")

for bad in ([], ):
    try:
        train_test_split_scratch(np.array(bad), np.array(bad))
    except ValueError as e:
        print("empty input →", e)

train 3500  test 1500
class balance — train 0.5003 0.4993  full 0.5000
empty input → cannot split an empty dataset

Stratification preserved to four decimals, and the empty case raises rather than returning nonsense. Saying “let me handle the empty case and a test_size outside (0,1)” costs two lines and is noticed.

What the round scores

Practice

nested python loops  41.284s
vectorised            0.004s

RuntimeWarning: overflow encountered in exp

my AUC 0.9401   sklearn AUC 0.9403   max |coef diff| 0.0287

seed 0-2: inertia 7305.41
seed 3:   inertia 9418.77

Next: classical ML depth — trees, ensembles, and the questions behind “which model”.