ML coding · Probabilistic models
Gaussian naive Bayes classifier
Fit a Gaussian naive Bayes model on X_train (n × d) with integer labels
y_train, and return the predicted label for each row of X_test.
- Prior of class c: the fraction of training rows with label c.
- For each class and feature: the mean and the biased variance (divide by the class count), plus ε = 10⁻⁹ × the largest variance of any feature over the whole training set.
- Predict the class with the largest log prior plus Σj log N(xj; μcj, σ²cj).
Every feature varies within every class in the tests.
Examples
Input X_train = [[0, 0], [1, 1], [0, 1], [5, 5], [6, 6], [5, 6]], y_train = [0, 0, 0, 1, 1, 1], X_test = [[0.5, 0.5], [5.5, 5.5]]
Output [0, 1]
Two well-separated clusters: each test point goes to the nearer one.
Input X_train = [[0], [2], [10], [12]], y_train = [0, 0, 1, 1], X_test = [[6.1]]
Output [1]
Equal variances and priors: the boundary is the midpoint, 6.
Input X_train = [[0], [1], [2], [0], [20]], y_train = [0, 0, 0, 1, 1], X_test = [[1]]
Output [0]
Class 1 is so spread out that its density at 1 is low: class 0 wins.
Constraints
- 2 ≤ n ≤ 5000, 1 ≤ d ≤ 50
- Every class has at least two training rows
- Labels are integers
Hints
Hint 1
np.unique(y) gives the classes. For each, take the mean and np.var of its rows.
Hint 2
The log density of a normal is −½ log(2πσ²) − (x − μ)²/(2σ²). Sum it over features.
Hint 3
Broadcast test rows against class means with T[:, None, :] − mu[None]: shape (n_test, n_classes, d). Sum over the last axis, add the log priors, take argmax over classes.
Where candidates go wrong
- Multiplying densities instead of adding log densities: with fifty features the products underflow and every class scores zero.
- Using the unbiased variance. Naive Bayes, like scikit-learn, uses the maximum-likelihood estimate.
- Forgetting the class prior, which matters whenever the classes are unbalanced.
- Returning the index of the class rather than the class label when labels are not 0, 1, 2, …
The idea
Fitting is a single pass of counting: the prior, mean and variance of every class. The posterior of class c for a point x is proportional to πc Πj N(xj; μcj, σ²cj), and its logarithm is a sum, which keeps everything finite and turns prediction into one broadcast expression over (test point, class, feature) followed by an argmax.
The ε added to every variance guards against zero variance, which would put an infinite density on a single value. Scaling it by the largest variance in the data makes it negligible whatever the units.
Time O(nd) to fit, O(n_test · C · d) to predict · Space O(Cd + n_test · C · d)
Reference solution
import numpy as np
def gaussian_nb(X_train, y_train, X_test):
X = np.asarray(X_train, dtype=float)
y = np.asarray(y_train)
T = np.asarray(X_test, dtype=float)
classes = np.unique(y)
eps = 1e-9 * X.var(axis=0).max() # scikit-learn's var_smoothing
mu = np.array([X[y == c].mean(axis=0) for c in classes]) # (C, d)
var = np.array([X[y == c].var(axis=0) for c in classes]) + eps
log_prior = np.log(np.array([np.mean(y == c) for c in classes]))
# log N(x; mu, var) summed over features, for every test row and class: (n_test, C)
ll = -0.5 * (np.log(2 * np.pi * var)[None] + (T[:, None, :] - mu[None]) ** 2 / var[None]).sum(axis=2)
return classes[np.argmax(ll + log_prior, axis=1)].tolist()Follow-ups an interviewer asks
- Return the posterior probabilities as well, normalised with log-sum-exp.
- Multinomial naive Bayes for word counts, with Laplace smoothing.
- Fit incrementally from batches: keep counts, sums and sums of squares per class.
- Why can naive Bayes classify well while its probabilities are badly calibrated?
Python runs in your browser with NumPy. The first run downloads the runtime, which takes a few seconds.
Why interviewers ask it
It is the simplest generative classifier, and writing it tests three things quickly: whether you work in log space, whether you can vectorise over test points, classes and features at once, and whether you know scikit-learn's variance-smoothing convention or can reason your way to it.
The theory behind it is in the lesson Naive Bayes, EM, Viterbi and Markov chains from scratch.
Keep going
Gaussian naive Bayes classifier
Fit class priors, means and variances, then predict by the largest log posterior, vectorised over test points and classes.
Open the problem → Medium20 minOne EM step for a Gaussian mixture
Compute responsibilities in log space, then refit weights, means and variances: one expectation-maximisation step for a 1-D Gaussian mixture, from scratch..
Open the problem → Medium25 minThe Viterbi algorithm for a hidden Markov model
Find the most likely sequence of hidden states in an HMM with dynamic programming in log space, with back-pointers to recover the path.
Open the problem → Easy12 minStationary distribution of a Markov chain
Solve πP = π with Σπ = 1 by replacing one redundant equation, in NumPy.
Open the problem →