Implement Naive Bayes classifier from scratch
Company: Pinterest
Role: Machine Learning Engineer
Category: Machine Learning
Difficulty: hard
Interview Round: Online Assessment
Implement a **Naive Bayes** classifier from scratch (you may use NumPy).
Write a class with:
- `fit(X, y)`: estimate class priors and feature likelihood parameters.
- `predict(X)`: compute posteriors (or unnormalized log-posteriors) and output the most likely class.
Specify which variant you implement (e.g., **Multinomial Naive Bayes** for count features, **Bernoulli NB** for binary features, or **Gaussian NB** for continuous features). Your implementation should clearly compute:
- prior \(P(y=c)\)
- likelihood parameters for \(P(x\mid y=c)\)
- posterior scoring using Bayes’ rule with numerical stability (typically log-probabilities).
Overview: This question evaluates understanding of probabilistic classification, specifically competence in estimating class priors, feature likelihood parameters, and computing posterior scores for Naive Bayes variants.
Read the full Pinterest Machine Learning Engineer interview experience this question came from
Community answers
Answer by swemleds
import numpy as np
class GaussianNaiveBayes:
def fit(self, X, y):
X = np.asarray(X, dtype=float)
y = np.asarray(y)
# 1. Find unique classes
self.classes = np.unique(y)
# 2. Initialize parameters for each class
self.priors = []
self.means = []
self.vars = []
for c in self.classes:
# 3. Select samples belonging to class c
X_c = X[y == c]
# 4. Compute prior: P(y = c)
self.priors.append(len(X_c) / len(X))
# 5. Compute mean of each feature for class c
self.means.append(np.mean(X_c, axis=0))
# 6. Compute variance of each feature for class c
# Add epsilon to avoid division by zero
self.vars.append(np.var(X_c, axis=0) + 1e-9)
self.priors = np.array(self.priors)
self.means = np.array(self.means)
self.vars = np.array(self.vars)
def predict(self, X):
X = np.atleast_2d(np.asarray(X, dtype=float))
predictions = []
# 1. Iterate through each test sample
for x in X:
scores = []
# 2. Compute posterior score for each class
for i, c in enumerate(self.classes):
# 3. Compute log prior: log P(y = c)
log_prior = np.log(self.priors[i])
# 4. Retrieve learned mean and variance
mean = self.means[i]
var = self.vars[i]
# 5. Compute Gaussian log likelihood:
# log P(x | y=c) = sum of feature log likelihoods
log_likelihood = -0.5 * np.sum(
np.log(2 np.pi var)
(x - mean) ** 2 / var
)
# 6. Compute unnormalized log posterior:
# log P(y=c) + log P(x | y=c)
scores.append(log_prior + log_likelihood)
# 7. Select class with highest posterior score
predictions.append(self.classes[np.argmax(scores)])
return np.array(