Train Your First Decision Tree

You've built scrapers, wrangled data, and called LLM APIs — but when it comes to explaining why a model made a decision, most black-box approaches leave you guessing. In production, that guesswork is a liability: regulators, stakeholders, and your own debugging sessions demand answers. Enter the decision tree, the most interpretable machine learning model you'll ever train — and in this lesson you'll go from zero to a working classifier that you can explain to a non-technical manager in plain English.

The problem this lesson solves

Modern AI feels magical, but magic doesn't ship well. When a customer is denied a loan or a medical diagnosis is flagged, you need to justify the model's behavior. Black-box models like deep neural networks can give you incredible accuracy but leave you blind to the reasoning behind individual predictions.

The pain is real: debugging a misbehaving model without insight is like fixing a car with the hood welded shut. You waste hours guessing which input mattered, whether the training data was skewed, or if there's a subtle bug in your feature pipeline.

A decision tree attacks this problem head-on. It's a transparent, rule-based model that literally draws its decision paths as a set of if-then-else rules. When you train your first decision tree, you're not just building a model — you're building a map of your data's logic that you can read, explain, and audit.

Core concept / mental model

Think of a decision tree as a 20 questions game played by a very patient coach. You start with the entire dataset at the root, then ask the single most informative question (e.g., "Is age > 30?"). Based on the answer, you split the data into two branches. Each branch poses another question, further narrowing down the possibilities, until you reach a leaf node — a final prediction.

The tree is built top-down, choosing at each node the feature and threshold that best separates the classes. The classic measure of "best" is information gain, which quantifies how much uncertainty (entropy) is reduced by a given split. The more homogeneous the resulting groups, the higher the gain.

Key terms to know

• Root node – the first split, the most important question.
• Internal nodes – subsequent questions.
• Leaf nodes – terminal predictions.
• Entropy – a measure of impurity; 0 means all same class, 1 means perfectly mixed.
• Information gain – reduction in entropy after a split.
• Pruning – removing branches that add little predictive power to prevent overfitting.

How it works step by step

Training a decision tree is deceptively simple, but each step matters:

1. Load and clean your data – remove missing values, encode categorical features, and split into features (X) and target (y).
2. Split into training and test sets – a typical 80/20 split, ideally stratified if your classes are imbalanced.
3. Initialize the tree – using scikit-learn's DecisionTreeClassifier.
4. Fit the model – call .fit(X_train, y_train); the algorithm greedily selects optimal splits.
5. Evaluate – score on the test set with accuracy, precision, recall, or a confusion matrix.
6. Inspect and visualize – export the tree as text or an image to understand its logic.
7. Tune hyperparameters – limit depth, minimum samples per leaf, etc., to avoid overfitting.

Cause and effect are crystal clear: each split reduces impurity, and the leaf you end up in determines the prediction. This is why the model is inherently explainable — you can trace any prediction back to a series of if-then rules.

Hands-on walkthrough

Let's train a decision tree on a classic dataset: iris flower classification. You'll need scikit-learn and pandas installed (pip install scikit-learn pandas).

1. Load the data and prepare it

from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
import pandas as pd

# Load dataset
iris = load_iris()
X = pd.DataFrame(iris.data, columns=iris.feature_names)
y = pd.Series(iris.target)

# Split into train/test
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)

print(f"Training samples: {len(X_train)}, Test samples: {len(X_test)}")

Expected output:

Training samples: 120, Test samples: 30

2. Train the decision tree

from sklearn.tree import DecisionTreeClassifier

# Create and train the model
clf = DecisionTreeClassifier(random_state=42)
clf.fit(X_train, y_train)

# Evaluate on test set
accuracy = clf.score(X_test, y_test)
print(f"Test accuracy: {accuracy:.2f}")

Expected output (may vary slightly):

Test accuracy: 1.00

On this toy dataset, an unconstrained tree often achieves perfect accuracy, but that's a red flag for overfitting in real-world data.

3. Visualize the tree

from sklearn.tree import export_text

# Print the rules as text
print(export_text(clf, feature_names=list(X.columns)))

Partial output:

|--- petal width (cm) <= 0.80
|   |--- class: 0
|--- petal width (cm) >  0.80
|   |--- petal width (cm) <= 1.75
|   |   |--- class: 1
|   |--- petal width (cm) >  1.75
|   |   |--- class: 2

The tree uses only one feature — petal width — because it alone separates the classes perfectly here. This is the interpretability payoff: the entire model boils down to a simple rule.

4. Make predictions and explain them

# Predict for a new sample
sample = [[5.1, 3.5, 1.4, 0.2]]  # sepal length, width, petal length, width
prediction = clf.predict(sample)
print(f"Predicted class: {iris.target_names[prediction[0]]}")

# Trace the decision path
print(clf.decision_path(sample).toarray())

Expected output:

Predicted class: setosa
[[True, True, False, False]]

The path array shows which nodes the sample passed through — every prediction can be audited.

Compare options / when to choose what

Decision trees are powerful, but they're not the only game in town. Here's how they stack up against common alternatives:

When to choose a decision tree

• You need explainability: loan approvals, medical diagnostics, compliance-heavy domains.
• You have tabular data with mixed feature types and don't want heavy preprocessing.
• You need a quick baseline to compare against more complex models.
• You have limited compute — trees train fast on CPU.

When to avoid it

• High-dimensional data (e.g., images, text) — other models will likely perform better.
• Large datasets — deep trees can become unwieldy and slow to predict.
• When you need top accuracy — a single tree often underperforms ensembles.

Variations worth knowing

• Random Forest – an ensemble of many trees, averaged to reduce overfitting.
• Gradient Boosted Trees (XGBoost, LightGBM) – sequential trees that correct errors, often state-of-the-art for tabular data.
• Pruning techniques – pre-pruning (limiting depth) and post-pruning (removing branches after training) to improve generalization.

Troubleshooting & edge cases

Overfitting — perfect training accuracy, poor test accuracy

This is the most common pitfall. Your tree memorized the training data instead of learning patterns.

Fix: Limit the tree's complexity:

clf = DecisionTreeClassifier(max_depth=3, min_samples_leaf=5, random_state=42)

Experiment with max_depth and min_samples_leaf using cross-validation.

Unbalanced classes

If one class dominates, the tree will be biased toward it.

Fix: Use class_weight='balanced' or resample your data with imblearn.

Categorical features with high cardinality

IDs or zip codes with thousands of unique values can cause the tree to split on meaningless noise.

Fix: Drop such features or group them into meaningful categories before training.

NaN values and error: Input contains NaN

scikit-learn's decision tree does not handle missing values natively.

Fix: Impute with SimpleImputer or drop rows with missing data.

Model makes no sense (e.g., predicts using an irrelevant feature)

Sometimes the tree finds spurious correlations in small datasets.

Fix: Use more training data, limit depth, and always validate with a holdout set.

RandomState warnings or non-reproducible results

If you don't set random_state, you'll get different trees each run.

Fix: Always specify random_state in your classifier for reproducibility.

What you learned & what's next

You've just trained your first decision tree end-to-end: you loaded a dataset, fit the model, evaluated its accuracy, and — crucially — read the exact rules it uses to make predictions. That's a huge step toward being an applied AI engineer who can ship models with confidence, not just accuracy.

You now know how to:

• Explain the core idea behind training a decision tree: recursive splitting to reduce impurity.
• Complete a practical exercise using scikit-learn, from preparation to prediction.
• Connect the interpretability benefit to real-world needs.

The next lesson in this track will build on this foundation — likely introducing ensemble methods like Random Forests, which combine many trees for even better performance while keeping some interpretability. You'll see how your single tree fits into a bigger predictive system.

Keep your tree — you'll want to compare its performance against an ensemble later. And remember: every time you train a model, ask yourself "Can I explain why this prediction happened?" If the answer is no, a decision tree might be your best friend.