ML Coding Interviews: What to Expect and How to Prepare

ML coding interviews are not like software engineering coding interviews. You will not be asked to implement a red-black tree. You will be asked to implement gradient descent, write a k-means algorithm, or build a neural network from scratch in NumPy. The skills that matter are different.

What You Will Encounter

Most ML roles have two types of coding rounds:

Type 1: Implement an ML algorithm from scratch

• Implement k-nearest neighbors
• Implement linear regression with gradient descent
• Implement k-means clustering
• Implement a decision tree (splitting criterion)
• Implement backpropagation for a 2-layer network

Type 2: ML data manipulation and feature engineering

• Given a dataset with issues (missing values, wrong types, outliers), clean it
• Compute rolling features over a time series
• Implement a train/val/test split that avoids data leakage
• Evaluate a classifier and explain the metrics

Type 3: Applied ML pipeline

• Given a description of a business problem, write a complete training pipeline
• Debug a broken training loop
• Optimize a slow prediction function

Implement Linear Regression from Scratch

This is the most common "implement from scratch" question. Know this cold:

pythonCopy
import numpy as np

class LinearRegression:
    def __init__(self, learning_rate=0.01, n_iterations=1000):
        self.learning_rate = learning_rate
        self.n_iterations = n_iterations
        self.weights = None
        self.bias = None
        self.loss_history = []

    def fit(self, X: np.ndarray, y: np.ndarray) -> 'LinearRegression':
        n_samples, n_features = X.shape
        self.weights = np.zeros(n_features)
        self.bias = 0

        for i in range(self.n_iterations):
            # Forward pass
            y_pred = self._predict(X)

            # Compute MSE loss
            loss = np.mean((y_pred - y) ** 2)
            self.loss_history.append(loss)

            # Compute gradients
            # dL/dw = (2/n) * X^T * (y_pred - y)
            dw = (2 / n_samples) * X.T @ (y_pred - y)
            db = (2 / n_samples) * np.sum(y_pred - y)

            # Update parameters
            self.weights -= self.learning_rate * dw
            self.bias -= self.learning_rate * db

        return self

    def _predict(self, X: np.ndarray) -> np.ndarray:
        return X @ self.weights + self.bias

    def predict(self, X: np.ndarray) -> np.ndarray:
        if self.weights is None:
            raise ValueError("Model not fitted yet")
        return self._predict(X)

    def score(self, X: np.ndarray, y: np.ndarray) -> float:
        """R² score"""
        y_pred = self.predict(X)
        ss_res = np.sum((y - y_pred) ** 2)
        ss_tot = np.sum((y - y.mean()) ** 2)
        return 1 - ss_res / ss_tot

Implement K-Means from Scratch

pythonCopy
import numpy as np

class KMeans:
    def __init__(self, k: int, max_iterations: int = 100, tol: float = 1e-4):
        self.k = k
        self.max_iterations = max_iterations
        self.tol = tol
        self.centroids = None
        self.labels_ = None

    def fit(self, X: np.ndarray) -> 'KMeans':
        n_samples = X.shape[0]

        # Initialize centroids: random sample from data (k-means++ is better)
        rng = np.random.default_rng(42)
        idx = rng.choice(n_samples, size=self.k, replace=False)
        self.centroids = X[idx].copy()

        for iteration in range(self.max_iterations):
            # Assign each point to nearest centroid
            distances = self._compute_distances(X)
            labels = np.argmin(distances, axis=1)

            # Recompute centroids
            new_centroids = np.array([
                X[labels == j].mean(axis=0) if (labels == j).any()
                else self.centroids[j]  # handle empty cluster
                for j in range(self.k)
            ])

            # Check convergence
            shift = np.max(np.linalg.norm(new_centroids - self.centroids, axis=1))
            self.centroids = new_centroids

            if shift < self.tol:
                print(f"Converged at iteration {iteration + 1}")
                break

        self.labels_ = labels
        return self

    def _compute_distances(self, X: np.ndarray) -> np.ndarray:
        """
        Compute distance from each point to each centroid.
        Returns: (n_samples, k) array
        """
        # Vectorized: avoid Python loop
        # ||x - c||² = ||x||² - 2x·c + ||c||²
        X_sq = np.sum(X ** 2, axis=1, keepdims=True)  # (n, 1)
        C_sq = np.sum(self.centroids ** 2, axis=1)      # (k,)
        cross = X @ self.centroids.T                    # (n, k)
        return X_sq + C_sq - 2 * cross

    def predict(self, X: np.ndarray) -> np.ndarray:
        return np.argmin(self._compute_distances(X), axis=1)

    def inertia(self, X: np.ndarray) -> float:
        """Sum of squared distances to nearest centroid."""
        labels = self.predict(X)
        return sum(
            np.sum((X[labels == j] - self.centroids[j]) ** 2)
            for j in range(self.k)
        )

Implement a Sigmoid Neural Network

pythonCopy
import numpy as np

class TwoLayerNetwork:
    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int,
                 learning_rate: float = 0.01):
        self.lr = learning_rate
        # He initialization for ReLU layers
        self.W1 = np.random.randn(input_dim, hidden_dim) * np.sqrt(2 / input_dim)
        self.b1 = np.zeros(hidden_dim)
        self.W2 = np.random.randn(hidden_dim, output_dim) * np.sqrt(2 / hidden_dim)
        self.b2 = np.zeros(output_dim)

    def sigmoid(self, x):
        return 1 / (1 + np.exp(-np.clip(x, -500, 500)))

    def relu(self, x):
        return np.maximum(0, x)

    def forward(self, X):
        self.X = X
        self.z1 = X @ self.W1 + self.b1
        self.a1 = self.relu(self.z1)
        self.z2 = self.a1 @ self.W2 + self.b2
        self.a2 = self.sigmoid(self.z2)
        return self.a2

    def backward(self, y):
        n = len(y)
        # Output layer gradient (binary cross-entropy + sigmoid combined)
        dz2 = self.a2 - y.reshape(-1, 1)          # (n, output_dim)
        dW2 = self.a1.T @ dz2 / n
        db2 = dz2.mean(axis=0)

        # Hidden layer gradient
        da1 = dz2 @ self.W2.T
        dz1 = da1 * (self.z1 > 0)                  # ReLU derivative
        dW1 = self.X.T @ dz1 / n
        db1 = dz1.mean(axis=0)

        # Update
        self.W1 -= self.lr * dW1
        self.b1 -= self.lr * db1
        self.W2 -= self.lr * dW2
        self.b2 -= self.lr * db2

    def fit(self, X, y, epochs=100):
        for epoch in range(epochs):
            output = self.forward(X)
            self.backward(y)
            if epoch % 20 == 0:
                loss = -np.mean(y * np.log(output + 1e-8) +
                                (1 - y) * np.log(1 - output + 1e-8))
                print(f"Epoch {epoch}: loss={loss:.4f}")

The Coding Interview Meta-Skills

Talk while you write. Interviewers are evaluating your thought process. Saying "I am going to initialize the centroids by random sampling because k-means++ would be more robust but takes more time to implement in an interview" shows more than just writing the random sampling.

Start with brute force. If you cannot see the optimal solution immediately, state the brute force approach first, then optimize. "The naive approach is O(n·k) per iteration which is fine for k-means in practice. If I needed to scale to very high k, I would use approximate nearest neighbor search."

Handle edge cases out loud. "What if a cluster becomes empty? I need to handle that - I will keep the old centroid so it does not crash."

Know your NumPy. Vectorized operations over loops. X @ W not np.dot(X, W) (same thing, cleaner). np.sum(axis=0) vs axis=1. Broadcasting rules.

What to Study

Prioritize in this order:

1. Linear algebra in NumPy: matrix multiply, broadcasting, vectorized operations
2. Linear regression and logistic regression from scratch
3. K-means from scratch
4. K-nearest neighbors from scratch
5. A simple neural network forward + backward pass
6. Pandas: groupby, merge, pivot, time series resampling, rolling features

Do not spend more than 20% of your prep time on LeetCode-style algorithm questions unless the role specifically asks for it. ML engineering roles care more about whether you can implement a training loop than whether you can solve a dynamic programming puzzle.

Common Mistakes

Reaching for nested loops when a better algorithm exists. Writing O(n²) solutions on sorted or indexed data is a common pattern interviewers flag immediately. If the input is sorted, binary search or a two-pointer approach almost always reduces the complexity, and demonstrating that awareness matters as much as getting a working solution.

Skipping edge cases in ML utility code. Not handling empty arrays, NaN inputs, or zero-variance features will fail on hidden test cases that well-designed interview graders include deliberately. A few lines of guard clauses at the top of your function show that you think about production robustness, not just the happy path.

Explaining the wrong dimension of the problem. Spending five minutes discussing model architecture when the interviewer asked about time complexity or data pipeline design signals a mismatch between what you know and what they are evaluating. Listen carefully for what dimension - correctness, efficiency, scalability, or modeling - the question is really targeting.