Train and Predict with Naive Bayes Practice Problem

This data science coding problem helps you practice Naive Bayes, train and predict with naive bayes, and implementation skills. Read the problem statement, write your solution, and strengthen your understanding of Naive Bayes.

• Problem ID: 197
• Problem key: 197-train-and-predict-with-naive-bayes
• URL: https://datacrack.app/solve/197-train-and-predict-with-naive-bayes
• Difficulty: hard
• Topic: Naive Bayes
• Module: Supervised Learning

Problem Statement

# 🧩 Train and Predict with Naive Bayes

---

### 🎯 Goal

Train a Naive Bayes model to learn reusable probabilities, then use the trained model to predict a new sample.

---

### 📖 Introduction

Naive Bayes has two separate phases.

**Training** receives features `X` and labels `y`. It learns and stores:

$$
P(class)
$$

and, for every feature column and every observed value:

$$
P(feature=value\mid class)
$$

**Prediction** receives the trained model and one new query. For every class, it starts with the stored prior and multiplies the stored likelihood for each query feature:

$$
Score(class)=P(class)\times\prod_{i=1}^{m}P(feature_i=value_i\mid class)
$$

Here, $m$ is the number of feature columns and $\prod$ means multiply one likelihood for each column. Naive Bayes makes a simplifying assumption that, once the class is known, each feature contributes its evidence independently. That assumption is why these likelihoods can be multiplied; it does not mean the features are always unrelated in real data. The class with the largest score is predicted. If scores tie, use the alphabetically smaller class.

The important boundary is this: prediction uses the probabilities already learned during training. It must not recalculate likelihoods by looking through `X` and `y` again.

---

### 📖 Training Data

| color | size | class |
|---|---|---|
| red | small | A |
| red | large | A |
| blue | small | B |
| blue | large | B |

---

### 💻 Task

Implement `naive_bayes_predict(X, y, query)`.

Your nested training helper should create a model with this structure:

```python
{
    "classes": [...],
    "priors": {class_label: ...},
    "likelihoods": {
        class_label: [
            {
                value_1: probability,
                value_2: probability,
                ...
            },  # feature column 0

            {
                value_1: probability,
                value_2: probability,
                ...
            },  # feature column 1

            ...
        ]
    }
}
```

Your function should:

- Store one prior for every class.
- Store a likelihood table for every class and feature column. Each table must contain every value observed in its column, including values with likelihood `0.0` for that class.
- Use `0.0` when a query value does not occur in a stored likelihood table.
- Return every prediction score rounded to 6 decimals and the winning class.
- Return `{"scores": {}, "prediction": None}` when the model has no classes.

---

### 📥 Input / 📤 Output

**Input:** categorical feature matrix `X`, labels `y`, and one categorical `query`.

**Prediction output:**

```python
{
    "scores": {"A": ..., "B": ...},
    "prediction": ...
}
```

---

### 🧩 Starter Code

```python
def naive_bayes_predict(X, y, query):
    def train_model():
        # Learn priors and likelihood tables.
        pass

    def predict_with_model(model):
        # Use only the learned model to calculate class scores.
        pass

    # Train once, then predict with that model.
    pass
```

---

### 💡 Example

```python
X = [
    ["red", "small"],
    ["red", "large"],
    ["blue", "small"],
    ["blue", "large"],
]
y = ["A", "A", "B", "B"]

naive_bayes_predict(X, y, ["red", "small"])
```

Expected output:

```python
{"scores": {"A": 0.25, "B": 0.0}, "prediction": "A"}
```

---

### ⚠️ Common Mistakes

- Recalculating likelihoods from `X` and `y` inside the nested prediction helper.
- Storing a likelihood table for only one feature column.
- Forgetting values that occur in another class, which should be stored with likelihood `0.0`.
- Choosing the smallest score or leaving ties undefined.