Handling Categorical Data with One-Hot Encoding: AI for Dummies #2

If you’ve been following along with this series, you’ll know that in the first post we built a simple ML model to predict diabetes. We kept things deliberately simple by using a dataset made up entirely of numerical data—things like blood pressure, glucose levels, and BMI.

But what happens when your data isn’t all numbers? What if you have a column that says "cat", "dog", and "fish"? Or one that says "red", "green", and "blue"?

Most machine learning models are, at their core, just a lot of maths. They deal in numbers. They multiply things, they add things, they calculate distances. They have absolutely no idea what to do with the word "red".

This is where One-Hot Encoding comes in. In this post, we’ll solve a real(ish) problem that forces us to deal with text categories, and we’ll see how to transform them into something our model can actually learn from.

The Problem: Predicting House Prices

We’re going to build a simple model to predict house prices. Our dataset will include numerical features like the number of bedrooms and the house’s square footage, but it will also include a categorical feature: the neighbourhood the house is in.

The neighbourhood is the tricky part. It has values like "Riverside", "Downtown", and "Suburbs". Our model can’t work with these as-is, so we need to encode them properly before we hand them over.

Our Setup

We’ll use Python with the usual suspects: pandas for data wrangling and scikit-learn for preprocessing and modelling. No Docker needed this time—we’ll keep it even simpler.

# requirements.txt
pandas
scikit-learn

Install them in a virtual environment:

python -m venv venv
source venv/bin/activate # or venv\Scripts\activate on Windows
pip install -r requirements.txt

The Practical Walkthrough

Here’s the complete script. It creates a dataset, applies One-Hot Encoding, trains a simple linear regression model, and evaluates it on a held-out test set.

# demo.py
import pandas as pd
from sklearn.preprocessing import OneHotEncoder
from sklearn.linear_model import LinearRegression
from sklearn.compose import ColumnTransformer
from sklearn.pipeline import Pipeline
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_absolute_error

# 1. Create our dataset (enough rows to actually split and learn from)
data = {
    'neighbourhood': [
        'Riverside', 'Downtown', 'Suburbs', 'Downtown', 'Riverside', 'Suburbs',
        'Riverside', 'Downtown', 'Suburbs', 'Riverside', 'Suburbs', 'Downtown'
    ],
    'bedrooms': [3, 2, 4, 1, 2, 3, 4, 3, 2, 1, 5, 2],
    'sqft':     [1500, 900, 2100, 600, 1100, 1800, 2200, 1600, 1000, 700, 2800, 950],
    'price':    [320000, 250000, 410000, 180000, 290000, 370000,
                 430000, 330000, 240000, 195000, 510000, 260000]
}
df = pd.DataFrame(data)

# 2. Split features (X) and target (y)
X = df[['neighbourhood', 'bedrooms', 'sqft']]
y = df['price']

# 3. Split into training and test sets before doing anything else
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=42)

# 4. Define the One-Hot Encoder for the categorical column
#   - drop='first' avoids the Dummy Variable Trap (more on this below)
#   - handle_unknown='ignore' prevents crashes if an unseen category appears at prediction time
preprocessor = ColumnTransformer(transformers=[
    ('onehot', OneHotEncoder(drop='first', handle_unknown='ignore'), ['neighbourhood'])
], remainder='passthrough')

# 5. Build a pipeline: preprocess, then train the model
pipeline = Pipeline(steps=[
    ('preprocessor', preprocessor),
    ('model', LinearRegression())
])

# 6. Train ONLY on the training set
pipeline.fit(X_train, y_train)

# 7. Evaluate on the unseen test set
y_pred = pipeline.predict(X_test)
mae = mean_absolute_error(y_test, y_pred)
print(f"Mean Absolute Error on test set: ${mae:,.0f}")

# 8. Make a prediction on a brand new house
new_house = pd.DataFrame([{
    'neighbourhood': 'Riverside',
    'bedrooms': 3,
    'sqft': 1400
}])
predicted_price = pipeline.predict(new_house)
print(f"Predicted price for new house: ${predicted_price[0]:,.0f}")

Running the script will output something like:

Mean Absolute Error on test set: $21,500
Predicted price for new house: $312,000

Nothing too fancy, but now we have an honest evaluation: the model’s predictions on houses it has never seen before are off by roughly $21,500 on average.

Code Explanation

1. Create the dataset: We build a pandas DataFrame with a mix of text and numerical columns. We use a reasonably sized dataset so the train/test split is actually meaningful.
2. Split features and target: X is our input (what we know), y is our output (what we want to predict).
3. Train/test split: We reserve 25% of the data as a held-out test set before any preprocessing happens. This is important—if you transform the data before splitting, information from the test set leaks into your training process.
4. Define the ColumnTransformer: We tell it to apply OneHotEncoder only to the neighbourhood column and pass everything else through unchanged. Notice the two extra parameters: drop='first' and handle_unknown='ignore'—we’ll explain these below.
5. Build the Pipeline: A Pipeline chains preprocessing and the model together into a single, clean object. The key benefit: the same transformations applied during training are automatically applied at prediction time—no risk of accidentally doing things differently.
6. Train on the training set: We call .fit() only on X_train. The pipeline encodes the data first, then trains the model on the result.
7. Evaluate on the test set: We run .predict() on X_test—data the model has never seen—and use Mean Absolute Error (MAE) to measure how far off our predictions are on average.
8. Predict a new house: We pass in a raw DataFrame with the neighbourhood as a plain string, and the pipeline handles encoding transparently.

Ok, But What Actually Is One-Hot Encoding?

Let’s slow down and think about why we can’t just use numbers directly.

A naive approach might be to encode the neighbourhoods as integers: Riverside = 1, Downtown = 2, Suburbs = 3. Looks clean, right? The problem is that the model will interpret those numbers as having a mathematical relationship. It’ll “think” that Suburbs (3) is three times more of something than Riverside (1), or that Downtown (2) is exactly halfway between them. That’s complete nonsense—there’s no such relationship between neighbourhoods.

A decent analogy here is flags. Imagine you need to represent your country’s flag on a form, and someone tells you to use the number 1 for Portugal, 2 for France, and 3 for Germany. If you run any maths on those numbers, you’d get gibberish. A flag is a flag—it’s a category, not a position on a scale.

One-Hot Encoding solves this by creating a new binary column for each unique category. Instead of one neighbourhood column with the value "Riverside", we end up with new columns:

is_Downtown	is_Riverside	is_Suburbs
0	1	0
1	0	0
0	0	1

Each row gets a 1 in exactly one column and 0 everywhere else. Now there’s no implied order or magnitude. The model sees three completely independent binary signals.

The Dummy Variable Trap (Why We Use drop='first')

Here’s something that catches a lot of people out, so it’s worth spending a moment on.

In the table above, notice that the three columns are perfectly predictable from each other. If is_Downtown = 0 and is_Riverside = 0, then is_Suburbs must be 1. Always. No exceptions. This means one of the three columns carries zero extra information—it is completely redundant.

When you feed perfectly redundant features into a Linear Regression model, it causes a problem called multicollinearity. The model’s coefficients become unstable and unreliable because the maths involved (matrix inversion) starts to break down. This is sometimes called the Dummy Variable Trap.

The fix is simple: always drop one of the encoded columns. That’s exactly what drop='first' does:

OneHotEncoder(drop='first', handle_unknown='ignore')

With 3 neighbourhoods, we end up with just 2 binary columns instead of 3—which is all the information we actually need. Note: if you were using a tree-based model (like Random Forest or XGBoost), you wouldn’t need to worry about this, since those models are immune to multicollinearity. But for Linear Regression, it matters.

A Quick Word on handle_unknown='ignore'

One last thing worth mentioning: what happens if, at prediction time, a neighbourhood name appears that the model has never seen during training? By default, OneHotEncoder will raise an error and crash. Not great for production.

Setting handle_unknown='ignore' tells the encoder to silently produce a row of all zeros for any unseen category. The model then makes a prediction based purely on the numerical features. It won’t be perfect, but it won’t blow up either—which is generally the more useful behaviour in the real world.

A Few Caveats Before You Go

We kept a few things simple here that are worth acknowledging.

First, our dataset is still tiny by real-world standards. A linear regression model with 5 effective features trained on a handful of rows doesn’t have much statistical room to breathe. The MAE number above should be taken as illustrative, not meaningful. In practice, you’d want hundreds or thousands of rows before trusting any model’s predictions.

Second, we glossed over feature explosion—the problem that occurs when you apply One-Hot Encoding to columns with hundreds or thousands of unique values (like city names or product SKUs). The resulting flood of binary columns makes the matrix very large and very sparse, which can slow down training significantly and hurt model performance. For those cases, there are better alternatives like Target Encoding or learned Embedding layers, which we’ll cover in a future article.

And that should be it. Cheers!