FastPrepGaussian Naive Bayes Classifier (Machine Learning OA)

Gaussian Naive Bayes Classifier (Machine Learning OA)

Pinterest logoPinterest● MediumFULLTIMEOA

Problem statement

Implement a Gaussian Naive Bayes classifier from scratch, without a machine-learning or numerical library. Train it on xTrain and yTrain, then return the predicted class label for each row of xTest, in the original row order.

Each row has m numeric features. yTrain[i] is the integer class label of training row xTrain[i]. The classes are the distinct values that appear in yTrain.

Training

For every class c, let n_c be the number of training rows labeled c, and let n be the total number of training rows. Compute:

  • The prior prior[c] = n_c / n.
  • For every feature j, the mean mean[c][j] of feature j over the rows of class c.
  • For every feature j, the variance var[c][j] = (1 / n_c) * sum((x[j] - mean[c][j])^2) + epsilon, where the sum runs over the rows x of class c. This divides by n_c, not n_c - 1. Adding epsilon keeps every variance positive, even for a class with one row or a constant feature.

Prediction

Naive Bayes treats the features as independent given the class, and models each feature with a normal distribution. For a test row x, the score of class c is its log posterior up to a shared constant:

score(c) = log(prior[c]) + sum over j of ( -0.5 * log(2 * pi * var[c][j]) - (x[j] - mean[c][j])^2 / (2 * var[c][j]) )

Here log is the natural logarithm. Predict the class with the largest score.

Practice contract

Work with log scores. Multiplying the densities directly can underflow to zero. The tests guarantee that, for every test row, the best class's score is larger than every other class's score by more than 10^-3, so small floating-point differences cannot change the answer. If yTrain has only one class, every test row is predicted as that class.

Function

predictGaussianNaiveBayes(xTrain: double[][], yTrain: int[], xTest: double[][], epsilon: double) → int[]

Examples

Example 1

xTrain = [[1.0],[2.0],[3.0],[6.0],[7.0],[8.0]]yTrain = [0,0,0,1,1,1]xTest = [[2.5],[6.5],[4.4]]epsilon = 0.001return = [0,1,0]

Class 0 has mean 2 and class 1 has mean 7. Both have prior 0.5 and variance 2/3 + 0.001 ≈ 0.6677. With equal priors and equal variances, the class whose mean is closer wins.

  • 2.5 is closer to 2, so it is class 0.
  • 6.5 is closer to 7, so it is class 1.
  • 4.4 has squared distance 5.76 to mean 2 and 6.76 to mean 7. The scores are about -5.72 for class 0 and -6.47 for class 1, so it is class 0.

Example 2

xTrain = [[0,0],[0.4,0.2],[-0.4,-0.2],[5,5],[1,9],[9,1]]yTrain = [2,2,2,7,7,7]xTest = [[1.5,1.5],[0,0.1],[4,4]]epsilon = 0.01return = [7,2,7]

Class 2 has means [0, 0] and small variances of about [0.1167, 0.0367]. Class 7 has means [5, 5] and variances of about [10.6767, 10.6767]. Both priors are 0.5.

  • [1.5, 1.5] is closer to [0, 0], but class 2 is so tight that this point is very unlikely under it. The scores are about -40.13 for class 2 and -6.05 for class 7, so it is class 7.
  • [0, 0.1] scores about 0.06 for class 2 and -7.19 for class 7, so it is class 2. A log score can be positive because a density can be larger than 1.
  • [4, 4] scores about -286.56 for class 2 and -4.99 for class 7, so it is class 7.

Example 3

xTrain = [[-1],[1],[-1],[1],[3],[5]]yTrain = [0,0,0,0,5,5]xTest = [[2.2],[2.6],[0]]epsilon = 0.5return = [0,5,0]

Class 0 has prior 2/3, mean 0, and variance 1 + 0.5 = 1.5. Class 5 has prior 1/3, mean 4, and the same variance 1.5. With equal variances, the score difference is log(2) ≈ 0.693 from the priors, plus the difference in (x - mean)^2 / 3.

  • For 2.2: 4.84 / 3 ≈ 1.613 for class 0 and 3.24 / 3 = 1.08 for class 5. Class 5 gains 0.533 from distance, but class 0 gains 0.693 from its prior. The scores are about -3.14 and -3.30, so it is class 0.
  • For 2.6: the distance terms are about 2.253 and 0.653. Class 5 gains 1.6, which beats the prior advantage. The scores are about -3.78 and -2.87, so it is class 5.
  • 0 is the mean of class 0, so it is class 0.

Constraints

  • The following bounds are practice assumptions.
  • 1 <= xTrain.length == yTrain.length <= 300
  • 1 <= xTest.length <= 300
  • 1 <= m <= 8, and every row of xTrain and xTest has exactly m values.
  • -100 <= xTrain[i][j], xTest[i][j] <= 100
  • 0 <= yTrain[i] <= 100
  • 0.001 <= epsilon <= 10
  • For every test row, the largest class score exceeds every other class score by more than 10^-3.

More Pinterest problems

See Pinterest hiring insights
Saved
public int[] predictGaussianNaiveBayes(double[][] xTrain, int[] yTrain, double[][] xTest, double epsilon) {
    // Write your code here.
}
Checking account…