Naïve Bayes
Java programming
machine learning
algorithm implementation
coding guidance

Implementing Naïve Bayes algorithm in Java - Need some guidance

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

Naive Bayes is a probabilistic classifier built on Bayes' theorem and a simplifying assumption: features are treated as conditionally independent given the class. Even though that assumption is rarely fully true, the algorithm works surprisingly well for text classification, spam detection, and other high-dimensional problems.

Focus on the multinomial version

For Java beginners, the multinomial Naive Bayes variant is a practical starting point because it fits bag-of-words text classification well. The model learns how often each token appears in each class and combines those token probabilities with the prior probability of each class.

At prediction time, the class with the highest posterior score wins.

Why log probabilities are used

Multiplying many small probabilities quickly underflows to zero in floating-point arithmetic. The usual fix is to sum logarithms instead of multiplying raw probabilities.

That changes this idea:

  • multiply priors and likelihoods

into this idea:

  • add log(prior) and log(likelihood) terms

The ranking of classes stays the same, but the computation is much more stable.

A small Java implementation

The example below trains a simple multinomial classifier with Laplace smoothing.

java
1import java.util.*;
2
3public class NaiveBayesClassifier {
4    private final Map<String, Map<String, Integer>> tokenCountsByClass = new HashMap<>();
5    private final Map<String, Integer> documentCounts = new HashMap<>();
6    private final Map<String, Integer> totalTokenCounts = new HashMap<>();
7    private final Set<String> vocabulary = new HashSet<>();
8    private int totalDocuments = 0;
9
10    public void train(String label, List<String> tokens) {
11        tokenCountsByClass.putIfAbsent(label, new HashMap<>());
12        documentCounts.put(label, documentCounts.getOrDefault(label, 0) + 1);
13        totalTokenCounts.put(label, totalTokenCounts.getOrDefault(label, 0) + tokens.size());
14        totalDocuments++;
15
16        Map<String, Integer> tokenCounts = tokenCountsByClass.get(label);
17        for (String token : tokens) {
18            vocabulary.add(token);
19            tokenCounts.put(token, tokenCounts.getOrDefault(token, 0) + 1);
20        }
21    }
22
23    public String predict(List<String> tokens) {
24        String bestLabel = null;
25        double bestScore = Double.NEGATIVE_INFINITY;
26
27        for (String label : documentCounts.keySet()) {
28            double score = Math.log((double) documentCounts.get(label) / totalDocuments);
29            Map<String, Integer> tokenCounts = tokenCountsByClass.get(label);
30            int totalTokensInClass = totalTokenCounts.get(label);
31
32            for (String token : tokens) {
33                int count = tokenCounts.getOrDefault(token, 0);
34                double probability = (count + 1.0) / (totalTokensInClass + vocabulary.size());
35                score += Math.log(probability);
36            }
37
38            if (score > bestScore) {
39                bestScore = score;
40                bestLabel = label;
41            }
42        }
43
44        return bestLabel;
45    }
46
47    public static void main(String[] args) {
48        NaiveBayesClassifier classifier = new NaiveBayesClassifier();
49
50        classifier.train("spam", List.of("win", "money", "now"));
51        classifier.train("spam", List.of("cheap", "money", "offer"));
52        classifier.train("ham", List.of("project", "meeting", "today"));
53        classifier.train("ham", List.of("lunch", "meeting", "schedule"));
54
55        String prediction = classifier.predict(List.of("cheap", "money"));
56        System.out.println(prediction);
57    }
58}

This model learns token counts per class, class document counts, and the global vocabulary size. When predicting, it scores each class and returns the label with the highest log-probability.

Laplace smoothing

Without smoothing, any token unseen in a class would make that class probability collapse to zero. Laplace smoothing avoids that by pretending every token has been seen once.

That is why the code uses count + 1.0 in the numerator and adds vocabulary.size() to the denominator.

Preparing the input data

Real classifiers usually need text preprocessing before training or prediction. Common steps include lowercasing, tokenization, punctuation removal, and optional stop-word filtering.

java
1import java.util.Arrays;
2import java.util.List;
3import java.util.stream.Collectors;
4
5public class TextPreprocessor {
6    public static List<String> tokenize(String text) {
7        return Arrays.stream(text.toLowerCase().split("\\W+"))
8                .filter(token -> !token.isBlank())
9                .collect(Collectors.toList());
10    }
11}

A model trained on one preprocessing strategy must use the same strategy during prediction. Otherwise, training and inference live in different token spaces and accuracy falls apart.

Common Pitfalls

A common mistake is multiplying raw probabilities directly instead of using logs. That often leads to underflow and confusing zero-like results.

Another issue is forgetting smoothing. Without it, unseen words punish a class too harshly and predictions become brittle.

It is also easy to build a correct implementation with poor data preparation. If training text is lowercased but prediction text is not, the token counts will not line up consistently.

Summary

  • Naive Bayes classifies by combining class priors with feature likelihoods.
  • The multinomial form is a strong starting point for text classification in Java.
  • Use log probabilities to avoid underflow during prediction.
  • Laplace smoothing prevents unseen tokens from zeroing out a class score.
  • Keep preprocessing consistent between training and prediction.

Related reading
Course
Intermediate
27 lessons
15 hours
DSA Fundamentals

Master algorithmic patterns and data structures through hands-on LeetCode-style problems - from arrays and hashing to dynamic programming and advanced graphs.

View the course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.