ML: Decision Trees in Python: Splitting, Pruning, Visualization

A python decision tree learns by asking simple yes/no questions, one after another, until it lands on an answer. This guide walks you through the whole thing in plain language: Gini impurity, information gain, how the tree picks each split, how to draw the tree out, how to control its depth so it does not overfit, and the moments when a tree beats a straight-line model.

“Trees provide the basis for some of the most effective prediction methods in use today.”

Leo Breiman, Statistical Modeling

Last Updated: July 2026 | Tested on: Python 3.14.6, scikit-learn 1.9.0 | Difficulty: Intermediate | Reading Time: 18 minutes

A decision tree makes a prediction by asking a chain of simple yes/no questions about your data. “Is income above 50,000?” Yes, go one way. “Is age above 30?” Yes, predict “approved.” Each question splits the data into cleaner groups, and the box at the end of the path holds the answer. This is the friendliest machine learning algorithm there is. You can draw the whole thing on paper and trace any input by hand.

Think of the game Twenty Questions, or that “Akinator” guessing game online. You think of something, the game asks “Is it alive?”, “Is it bigger than a car?”, “Does it fly?”, and each answer narrows the possibilities until only one thing is left. A decision tree learns that exact set of questions automatically, and it learns them in the smartest order so it reaches the answer fast. The only real skill is picking a good question at each step.

So how does the tree pick that question? At every step it tries each possible split and measures how cleanly it separates the classes. Gini impurity and information gain are the two yardsticks people use for that. A great split puts almost all the approved cases on one side and almost all the rejected cases on the other. A useless split leaves both sides just as mixed as before. The tree is greedy: at each node it grabs the single best split it can find right now and moves on.

Here is why people actually reach for trees at work. Say a developer named Viraj, 27, was building a loan approval model at a bank. His compliance team demanded a reason for every rejection. A logistic regression could only say “the probability came out to 0.38,” which is not a reason anyone can act on. The tree, on the other hand, said “Rejected because credit_score was below 650 and debt_ratio was above 0.40.” The team could read that, audit it, and explain it to a customer. That kind of plain-English trail is why banks, insurers, and hospitals often pick decision trees over fancier models.

Prerequisites

How Trees Split: Gini Impurity

Yes (55)No (45)Yes (40)No (15)Yes (20)No (25)Yes (8)No (12)All Data100 samples48 approved, 52 rejectedcredit_score> 650?debt_ratio<= 0.40?income> 40K?Approved35/40Rejected12/15debt_ratio<= 0.35?Rejected23/25Approved6/8Rejected10/12Python Decision Tree: Splitting Loan Data on Credit Score, Debt Ratio and Income

Tip: click the diagram to open it full screen, then use the zoom and pan controls for a closer look.

The diagram shows a loan-approval decision tree splitting the data through a chain of yes/no questions. At each node the algorithm picks the feature and the threshold that best separates approved from rejected, then sends the data left or right. The tree keeps growing until it reaches the leaf boxes that hold the final call. Read the numbers under each leaf (for example “Approved 35/40”) as “35 of the 40 people who landed here were actually approved.” This whole structure is easy to read: you can follow any single applicant from the top box down to a leaf and see exactly which conditions decided their fate.

The one thing you have to manage is depth, because a tree left to grow forever will memorize the training data instead of learning from it. That is what the pruning and max_depth settings are for.

Before any code, here is the plain idea of Gini impurity. Picture a bag of marbles. If every marble is the same color, the bag is “pure” and Gini is 0. If the bag is a perfect 50/50 mix of two colors, it is as messy as it gets and Gini hits 0.5 (for two classes). Gini just measures the chance that if you grabbed two marbles at random, they would be different colors. A good split is one that hands you two bags that are each closer to single-color than the bag you started with. Let us compute it by hand for a few small bags.

📄 gini_impurity.py: understanding the splitting criterion

import numpy as np

def gini(labels):
    """Gini impurity: 1 - sum(p_i^2) for each class."""
    if len(labels) == 0:
        return 0
    classes, counts = np.unique(labels, return_counts=True)
    probs = counts / len(labels)
    return 1 - np.sum(probs ** 2)

# Pure node (all same class) gives Gini = 0
pure = [1, 1, 1, 1, 1]
print(f"Pure [1,1,1,1,1]:         Gini = {gini(pure):.3f}")

# Maximally impure (50/50) gives Gini = 0.5
mixed = [0, 0, 0, 1, 1, 1]
print(f"Mixed [0,0,0,1,1,1]:      Gini = {gini(mixed):.3f}")

# Slightly impure
mostly_one = [1, 1, 1, 1, 0]
print(f"Mostly-one [1,1,1,1,0]:   Gini = {gini(mostly_one):.3f}")

# Evaluate a split: income > 50K?
# Before split: 10 approved, 10 rejected
parent = [1]*10 + [0]*10
print(f"\nParent (10 approved, 10 rejected): Gini = {gini(parent):.3f}")

# Split on income > 50K. left: 8 approved, 2 rejected. right: 2 approved, 8 rejected
left = [1]*8 + [0]*2
right = [1]*2 + [0]*8
weighted_gini = (len(left)/len(parent)) * gini(left) + \
                (len(right)/len(parent)) * gini(right)
info_gain = gini(parent) - weighted_gini
print(f"Left (income>50K):  Gini = {gini(left):.3f}")
print(f"Right (income<=50K): Gini = {gini(right):.3f}")
print(f"Weighted child Gini: {weighted_gini:.3f}")
print(f"Information gain:    {info_gain:.3f}")

▶ Output

Pure [1,1,1,1,1]:         Gini = 0.000
Mixed [0,0,0,1,1,1]:      Gini = 0.500
Mostly-one [1,1,1,1,0]:   Gini = 0.320

Parent (10 approved, 10 rejected): Gini = 0.500
Left (income>50K):  Gini = 0.320
Right (income<=50K): Gini = 0.320
Weighted child Gini: 0.320
Information gain:    0.180

What happened here: The numbers line up with the marble picture. A pure bag scores 0, a 50/50 bag scores 0.500, and a "mostly one color" bag (four ones and a zero) sits in between at 0.320. The split is the interesting part. We started with a fully mixed parent (Gini 0.500) and split it into two cleaner groups, each at Gini 0.320. The weighted average of the children, 0.320, is lower than the parent, and the drop of 0.180 is the information gain. A tree tries every possible split and keeps the one with the biggest information gain, because the biggest drop means the cleanest separation. That single rule, repeated at every node, is the entire learning algorithm.

Building a Python Decision Tree with scikit-learn

Computing Gini by hand is great for understanding, but you will never do it in real life. scikit-learn does all of that searching for you. It is like switching from long division on paper to punching the numbers into a calculator: the method is identical, the machine just does the grinding. Below we build a small loan-approval dataset where the real rule is "approve if income is above 40K and credit score is above 650 and debt ratio is below 0.4," then sprinkle in some noise so it is not too clean, and let the Python decision tree rediscover the pattern on its own.

📄 decision_tree.py: classification tree with text rules

import numpy as np
from sklearn.tree import DecisionTreeClassifier, export_text
from sklearn.model_selection import train_test_split, cross_val_score
from sklearn.metrics import accuracy_score

rng = np.random.default_rng(42)
n = 500

# Loan approval dataset
income = rng.integers(20000, 150000, n).astype(float)
credit_score = rng.integers(300, 850, n).astype(float)
employment_years = rng.integers(0, 20, n).astype(float)
debt_ratio = rng.uniform(0, 0.8, n)

# Approval logic: income > 40K AND credit > 650 AND debt < 0.4
approved = ((income > 40000) & (credit_score > 650) & (debt_ratio < 0.4)).astype(int)
# Add some noise
noise_idx = rng.choice(n, 40, replace=False)
approved[noise_idx] = 1 - approved[noise_idx]

X = np.column_stack([income, credit_score, employment_years, debt_ratio])
feature_names = ["income", "credit_score", "employment_years", "debt_ratio"]

X_train, X_test, y_train, y_test = train_test_split(X, approved, test_size=0.2,
                                                       random_state=42, stratify=approved)

# Train with controlled depth to prevent overfitting
tree = DecisionTreeClassifier(max_depth=4, min_samples_leaf=10, random_state=42)
tree.fit(X_train, y_train)

print(f"Train accuracy: {tree.score(X_train, y_train):.3f}")
print(f"Test accuracy:  {tree.score(X_test, y_test):.3f}")
print(f"Tree depth: {tree.get_depth()}")
print(f"Number of leaves: {tree.get_n_leaves()}")

# Feature importance
print(f"\nFeature Importance:")
for name, imp in sorted(zip(feature_names, tree.feature_importances_), key=lambda x: -x[1]):
    print(f"  {name}: {imp:.3f}")

# Text visualization of the tree
print(f"\nTree Rules:")
print(export_text(tree, feature_names=feature_names, max_depth=3))

▶ Output

Train accuracy: 0.917
Test accuracy:  0.910
Tree depth: 4
Number of leaves: 13

Feature Importance:
  debt_ratio: 0.511
  credit_score: 0.329
  income: 0.147
  employment_years: 0.014

Tree Rules:
|--- credit_score <= 657.50
|   |--- employment_years <= 8.50
|   |   |--- credit_score <= 536.50
|   |   |   |--- income <= 56852.00
|   |   |   |   |--- class: 0
|   |   |   |--- income >  56852.00
|   |   |   |   |--- class: 0
|   |   |--- credit_score >  536.50
|   |   |   |--- income <= 61050.00
|   |   |   |   |--- class: 0
|   |   |   |--- income >  61050.00
|   |   |   |   |--- class: 0
|   |--- employment_years >  8.50
|   |   |--- credit_score <= 320.50
|   |   |   |--- class: 0
|   |   |--- credit_score >  320.50
|   |   |   |--- credit_score <= 628.50
|   |   |   |   |--- class: 0
|   |   |   |--- credit_score >  628.50
|   |   |   |   |--- class: 0
|--- credit_score >  657.50
|   |--- debt_ratio <= 0.39
|   |   |--- income <= 50973.00
|   |   |   |--- class: 0
|   |   |--- income >  50973.00
|   |   |   |--- debt_ratio <= 0.11
|   |   |   |   |--- class: 1
|   |   |   |--- debt_ratio >  0.11
|   |   |   |   |--- class: 1
|   |--- debt_ratio >  0.39
|   |   |--- credit_score <= 762.50
|   |   |   |--- credit_score <= 736.00
|   |   |   |   |--- class: 0
|   |   |   |--- credit_score >  736.00
|   |   |   |   |--- class: 0
|   |   |--- credit_score >  762.50
|   |   |   |--- class: 0

What happened here: The tree got 91.7% on the training data and held onto 91.0% on data it had never seen, so it learned a real pattern rather than memorizing. Look at the feature importances: debt_ratio came out on top at 0.511, credit_score next at 0.329, then income at 0.147, while employment_years barely registered at 0.014. The tree figured out on its own which columns matter and which are noise, and it agrees with the rule we baked in (employment years was never part of the real approval logic, so it sits near zero).

Now read the text rules like a flowchart. Follow the credit_score > 657.50 branch (the second top-level branch in the printout) all the way down: if credit_score is above 657.50, then debt_ratio is at or below 0.39, and income is above 50,973, the tree predicts class 1, which is "approved." That is a sentence a compliance officer can read out loud. Notice the flip side too: almost every other path ends in class 0, because in this dataset only about 23% of applicants were approved, so "reject" is the safe default and the tree only says "approve" when several conditions line up.

The repeated splits that all land on class 0 (like the two income branches that both predict 0) are the tree being cautious near a boundary, which is exactly the kind of thing pruning cleans up next.

Pruning: Controlling Overfitting

A tree with no limits is like a student who memorizes the answer key word for word. Ace the practice exam, flunk the real one. Pruning is the fix: you stop the tree from growing too deep so it learns the general pattern instead of memorizing every single training row. The clearest way to see this is to grow trees at different depths and watch the gap between training accuracy and cross-validation accuracy. Cross-validation hides part of the data, trains on the rest, and checks the hidden part, so it is a fair stand-in for "data the model has never seen." That train-to-validation gap is the first thing to check on any Python decision tree you grow.

📄 pruning.py: the effect of max_depth on overfitting

import numpy as np
from sklearn.tree import DecisionTreeClassifier
from sklearn.model_selection import cross_val_score
from sklearn.datasets import make_classification

X, y = make_classification(n_samples=500, n_features=10, n_informative=5,
                            random_state=42)

print(f"{'max_depth':>10} | {'Train Acc':>10} | {'CV Acc':>10} | {'Gap':>8} | {'Leaves':>6}")
print("-" * 55)

for depth in [2, 3, 5, 10, None]:
    tree = DecisionTreeClassifier(max_depth=depth, random_state=42)
    tree.fit(X, y)
    train_acc = tree.score(X, y)
    cv_scores = cross_val_score(tree, X, y, cv=5)
    cv_acc = cv_scores.mean()
    gap = train_acc - cv_acc
    leaves = tree.get_n_leaves()
    d = str(depth) if depth else "None"
    print(f"{d:>10} | {train_acc:>9.3f} | {cv_acc:>9.3f} | {gap:>7.3f} | {leaves:>6}")

▶ Output

 max_depth |  Train Acc |     CV Acc |      Gap | Leaves
-------------------------------------------------------
         2 |     0.792 |     0.766 |   0.026 |      4
         3 |     0.876 |     0.828 |   0.048 |      8
         5 |     0.970 |     0.866 |   0.104 |     25
        10 |     1.000 |     0.862 |   0.138 |     42
      None |     1.000 |     0.862 |   0.138 |     42

What happened here: Read down the table and the story jumps out. At depth 10 and at no limit at all, the tree scores a perfect 1.000 on training data but only 0.862 on data it has not seen. That gap of 0.138 is the overfitting alarm bell: the tree memorized noise. At depth 2 it is too simple and underfits, scoring just 0.792 on training. Depth 5 lands the best cross-validation score here at 0.866 with a much smaller gap, so it is the sweet spot for this dataset.

One more thing worth noticing: depth 10 and "None" produced identical rows (42 leaves, same scores), because on this particular data the tree naturally stops growing around depth 10 even when you allow it to go further. The takeaway is steady across every dataset: a bigger gap between train and cross-validation accuracy, and more leaves, both mean more memorizing. Cap the depth or set min_samples_leaf and the problem goes away.

Decision Trees for Regression

A Python decision tree is not just for yes/no answers. Swap the classifier for a regressor and the same idea predicts a number, like a house price or a temperature. The split logic still works the same way, except instead of measuring class purity it measures how spread out the numbers are in each group, and each leaf simply predicts the average of the training values that landed in it. The fun part: a tree handles curvy, non-linear data without you adding a single polynomial feature. Here we fit a noisy sine wave, which a straight line could never follow.

📄 regression_tree.py: predicting continuous values

import numpy as np
from sklearn.tree import DecisionTreeRegressor
from sklearn.model_selection import cross_val_score

rng = np.random.default_rng(42)
X = rng.uniform(0, 10, 200).reshape(-1, 1)
y = np.sin(X.ravel()) * 10 + rng.normal(0, 1, 200)

for depth in [2, 5, 10]:
    tree = DecisionTreeRegressor(max_depth=depth, random_state=42)
    scores = cross_val_score(tree, X, y, cv=5, scoring="r2")
    print(f"Depth {depth:>2}: R² = {scores.mean():.3f} ± {scores.std():.3f}")

# Trees handle non-linear data naturally, no polynomial features needed!
print("\nDecision trees capture non-linear patterns without feature engineering.")
print("Each leaf predicts the mean of training samples that reached it.")

▶ Output

Depth  2: R² = 0.708 ± 0.026
Depth  5: R² = 0.953 ± 0.013
Depth 10: R² = 0.956 ± 0.008

Decision trees capture non-linear patterns without feature engineering.
Each leaf predicts the mean of training samples that reached it.

What happened here: R² is just "how much of the wiggle in the data did we capture," where 1.0 is perfect and 0 is no better than guessing the average. A shallow depth-2 tree only manages 0.708, because four leaves cannot trace a full sine curve. Bump it to depth 5 and the score jumps to 0.953, and depth 10 nudges a hair higher to 0.956. The little number after the plus-or-minus is the spread across the five cross-validation folds, so smaller is steadier: depth 10 is both the highest scoring and the most consistent here.

On this clean, low-noise sine wave the deeper tree is not overfitting, it genuinely fits the curve better. That will not always be true, so you still check the train-versus-validation gap like we did in the pruning table. The headline stays the same: the tree carved a smooth-looking curve out of plain yes/no splits, with zero feature engineering.

Common Mistakes

The single biggest mistake with trees is leaving them unpruned. A fresh DecisionTreeClassifier() with no limits keeps splitting until every leaf is pure, which often means one leaf per training row. That is not learning, it is memorizing, and it falls apart on new data. Always pass at least one guardrail.

❌ Mistake: not setting max_depth or min_samples_leaf

# An unpruned decision tree WILL overfit. Every time.
# It will create one leaf per training sample if allowed.

# Good defaults:
# max_depth=5-10 (rarely need more)
# min_samples_leaf=5-20 (prevents tiny leaves)
# min_samples_split=10-20 (prevents splitting tiny nodes)

# Or use cost-complexity pruning:
# tree = DecisionTreeClassifier(ccp_alpha=0.01)
# Higher ccp_alpha = more pruning

print("Always prune. An unpruned tree is like a student who memorizes")
print("the textbook word-for-word but can't answer a new question.")

▶ Output

Always prune. An unpruned tree is like a student who memorizes
the textbook word-for-word but can't answer a new question.

What happened here: Pick whichever guardrail fits your problem and you are safe. max_depth caps how many questions deep the tree can go. min_samples_leaf refuses to make a leaf out of just a handful of rows, so the tree cannot carve out a tiny box around one weird sample. ccp_alpha (cost-complexity pruning) grows the full tree, then trims back the branches that barely help, with a higher value meaning a more aggressive trim. In day-to-day work, start with max_depth=5 and min_samples_leaf=10, then tune from there. Never ship a tree with no limits at all.

Practice Exercises

  1. Exercise 1: Extend the hand-written gini() function to also compute entropy (the formula is the negative sum of p times log2(p) over the classes). Run both on the same bags of labels and confirm they agree on which split is best, even though the raw numbers differ.
  2. Exercise 2: Load the built-in load_iris dataset from scikit-learn, train a DecisionTreeClassifier(max_depth=3), and print the rules with export_text. Trace one flower from the top of the tree to its leaf by hand and check the prediction against the model.
  3. Exercise 3: Take the pruning loop from this post and add a column for ccp_alpha values like 0.0, 0.005, 0.01, and 0.02. Find the alpha that gives the smallest gap between train and cross-validation accuracy, then compare that tree's leaf count to the unpruned one.

Conclusion

You now have the whole picture of a Python decision tree. You saw how a tree asks a chain of yes/no questions, how Gini impurity and information gain let it pick the best question at each step, and how to build one in a few lines with scikit-learn and read its rules out loud. You also learned the one habit that keeps trees honest: prune them with max_depth, min_samples_leaf, or ccp_alpha so they learn the pattern instead of memorizing the data. And you saw the same machinery predict continuous numbers with a regression tree, tracing a curve that a straight line never could.

The natural next step is the random forest tutorial, which fixes the one big weakness of a single tree: instability. A forest trains many trees on different slices of the data and averages them, trading a little readability for a lot of accuracy. Work through the exercises above first, then move on. For the full path from basics to advanced models, visit the Python + AI/ML tutorial series home.

Frequently Asked Questions

Should I use Gini impurity or entropy?

In practice, they produce nearly identical trees. Gini is slightly faster to compute (no logarithm). scikit-learn defaults to Gini. Use entropy only if you have a specific reason. Most practitioners never change this setting.

Do decision trees need feature scaling?

No! This is a major advantage. A Python decision tree splits on feature values directly (is income > 50K?) so the scale does not matter. You can mix features with different units without preprocessing.

Why are decision trees unstable?

Small changes in training data can produce completely different trees. This instability is their main weakness. Random Forests (random forest tutorial) solve this by training many trees on different data subsets and averaging their predictions.

Can decision trees handle missing values?

scikit-learn trees cannot handle NaN values directly, so you must impute (fill in) the gaps first. However, some implementations like XGBoost and LightGBM handle missing values natively by learning a default direction at each split.

Interview Questions on Decision Trees

Scenario questions, not trivia: this is the form this topic takes in a real interview.

Q: How does a decision tree decide which feature and threshold to split on at each node?

At every node it tries every feature and every candidate threshold, measures the impurity of the two resulting groups (Gini or entropy for classification, variance or MSE for regression), and keeps the split with the largest information gain, meaning the biggest drop from the parent's impurity to the weighted average of the children. It is greedy: it takes the single best split available right now without looking ahead. That one rule, repeated at every node, is the entire training algorithm.

Q: What is the difference between how a classification tree and a regression tree split and predict?

A classification tree measures class purity with Gini or entropy and each leaf predicts the majority class of the training samples that reached it. A regression tree measures how spread out the numbers are (variance or mean squared error) and each leaf predicts the average of the target values that landed in it. The tree structure and greedy split search are identical; only the impurity metric and the leaf output change. In scikit-learn you swap DecisionTreeClassifier for DecisionTreeRegressor.

Q: What does feature_importances_ actually measure, and what is one caveat?

It reports how much each feature reduced impurity across all the splits that used it, normalized to sum to 1, so a higher value means the feature did more of the separating work. The main caveat is that this impurity-based importance is biased toward high-cardinality features, like a near-unique ID column, because such columns give the tree many places to split. When you need a more trustworthy ranking, prefer permutation importance on held-out data.

Q: Your decision tree scores 100% on the training set but only 70% on the test set. What is happening and what do you change first?

That gap is textbook overfitting: with no limits the tree grew until nearly every leaf memorized a single training row, including its noise. The first thing to change is a depth or size guardrail, for example max_depth=5 and min_samples_leaf=10, then use cross-validation to find the setting where the train-versus-validation gap is smallest. Cost-complexity pruning with ccp_alpha is a good alternative that trims the least useful branches after growing the full tree.

Q: A stakeholder asks exactly why the model rejected one specific loan application. How do you produce that answer from a trained tree?

This is where trees shine over most models. Call export_text(tree, feature_names=...) to dump the full rule set, or use tree.decision_path(X) to get the exact sequence of nodes that one application passed through. You can then read the conditions along that path in plain English, such as "rejected because credit_score was below 657.50 and debt_ratio was above 0.39." That traceable, auditable reason is why regulated fields like banking and healthcare favor trees.

Q: Your training data is 95% "reject" and 5% "approve," and the tree predicts "reject" for almost everything. How do you fix it?

With a heavy class imbalance, always predicting the majority class already gives high accuracy, so the tree has little incentive to learn the rare class. Set class_weight="balanced" so mistakes on the minority class are penalized more, and stop judging the model by raw accuracy: look at precision, recall, and the confusion matrix for the "approve" class instead. Resampling techniques like SMOTE or simply collecting more positive examples also help.

Q: What is cost-complexity pruning and how does ccp_alpha control it?

Cost-complexity pruning grows the full tree, then works backward and removes the branches whose accuracy contribution does not justify their added complexity. The ccp_alpha parameter sets the price of each extra leaf: a higher alpha makes complexity more expensive, so more branches get trimmed and the tree ends up smaller and more general. You typically test a range of alpha values with cross-validation and pick the one with the best validation score.

Series: Python + AI/ML Cookbook, Part 5: Machine Learning

Go deeper: when you outgrow this post, scikit-learn documentation is the next stop.

Previous: ML: Confusion Matrix & Receiver Operating Characteristic (ROC) Curves

Next: ML: Random Forest, Bagging and Feature Importance

Series Home: Python + AI/ML Tutorial Series

RahulAuthor posts

Avatar for Rahul

Rahul is a passionate IT professional who loves to sharing his knowledge with others and inspiring them to expand their technical knowledge. Rahul's current objective is to write informative and easy-to-understand articles to help people avoid day-to-day technical issues altogether. Follow Rahul's blog to stay informed on the latest trends in IT and gain insights into how to tackle complex technical issues. Whether you're a beginner or an expert in the field, Rahul's articles are sure to leave you feeling inspired and informed.

No comment

Leave a Reply

Your email address will not be published. Required fields are marked *