major: changing the logic

Author: Rosander0

src/main/java/com/thealgorithms/machinelearning/MultinomialNaiveBayesClassifier.java

@@ -76,13 +76,14 @@ public void fit(double[][] features, int[] labels) {
         }
 
         int totalSamples = features.length;
-        for (Map.Entry<Integer, Integer> entry : classCounts.entrySet()) {
+        for (Map.Entry<Integer, double[]> entry : featureSums.entrySet()) {
             int label = entry.getKey();
-            int count = entry.getValue();
+            double[] sums = entry.getValue();
+            int count = classCounts.getOrDefault(label, 0);
+            double total = totalFeatureCount.getOrDefault(label, 0.0);
+
             logPriors.put(label, Math.log((double) count / totalSamples));
 
-            double[] sums = featureSums.get(label);
-            double total = totalFeatureCount.get(label);
             double denom = total + alpha * numFeatures;
             double[] logLikelihood = new double[numFeatures];
             for (int j = 0; j < numFeatures; j++) {
@@ -109,10 +110,10 @@ public int predict(double[] sample) {
         int bestLabel = -1;
         double bestScore = Double.NEGATIVE_INFINITY;
 
-        for (Map.Entry<Integer, Double> entry : logPriors.entrySet()) {
+        for (Map.Entry<Integer, double[]> entry : logLikelihoods.entrySet()) {
             int label = entry.getKey();
-            double score = entry.getValue();
-            double[] logLikelihood = logLikelihoods.get(label);
+            double[] logLikelihood = entry.getValue();
+            double score = logPriors.getOrDefault(label, Double.NEGATIVE_INFINITY);
             for (int j = 0; j < numFeatures; j++) {
                 score += sample[j] * logLikelihood[j];
             }

src/test/java/com/thealgorithms/machinelearning/MultinomialNaiveBayesClassifierTest.java

@@ -3,6 +3,7 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+
 import org.junit.jupiter.api.Test;
 
 class MultinomialNaiveBayesClassifierTest {