Decision Trees in Python: Build, Visualize, and Prune

Contents

Figure 1: A tree is a flowchart that learned its questions from data

Impurity and splits, a readable Iris tree, overfitting on purpose, then both kinds of pruning

A decision tree is the closest thing machine learning has to a flowchart you can read out loud. "Is petal length under 2.5 cm? Then it's a setosa."

That readability is why trees remain popular, and it's also why they come with a classic trap: left alone, a tree will grow until it memorizes your training data, producing a spectacular training score and a disappointing real-world one. This guide builds a tree with scikit-learn, shows you how to visualize it, demonstrates overfitting on purpose, and then fixes it with both kinds of pruning.

How a Tree Decides

A decision tree learns by asking a sequence of yes/no questions about your features. At each node it searches for the single split, such as "feature 3 ≤ 0.7", that best separates the classes. "Best" means the largest drop in impurity — a measure of how mixed the classes are in a node:

• Gini impurity (the default) is the chance of mislabeling a randomly chosen sample if you labeled it by the node's class mix.
• Entropy measures disorder in an information-theory sense; scikit-learn also offers log_loss.

In practice the two criteria usually produce very similar trees, so don't agonize over the choice. The algorithm is greedy: it picks the best split at each step without looking ahead, then recurses on each side until it hits a stopping rule or the nodes are pure. A prediction is just the majority class (or the average value, for regression trees) in the leaf a sample lands in.

Build and Visualize a Small Tree

We'll start with the Iris dataset because a shallow tree on it is small enough to read:

import matplotlib.pyplot as plt
from sklearn.datasets import load_iris
from sklearn.tree import DecisionTreeClassifier, export_text, plot_tree

iris = load_iris()
X, y = iris.data, iris.target

tree = DecisionTreeClassifier(max_depth=3, random_state=0).fit(X, y)

print(export_text(tree, feature_names=list(iris.feature_names)))

plt.figure(figsize=(12, 7))
plot_tree(tree, feature_names=iris.feature_names, class_names=iris.target_names,
          filled=True, rounded=True)
plt.show()

(We fit on all the data here purely for visualization. Never evaluate a model on the data it was trained on.)

Each box in the plot shows the split condition, the node's impurity (gini), the number of samples, the value (how many samples of each class), and the majority class. Color intensity reflects purity. export_text prints the same logic as indented rules, which is handy for logs, code reviews, and emailing a model to a skeptical colleague.

For publication-quality figures, export_graphviz writes a Graphviz file you can render with the Graphviz tools, and third-party packages such as dtreeviz offer richer charts. Check their current documentation before depending on them.

Watch a Tree Overfit

Now a harder problem. This synthetic dataset includes 10% label noise, so a perfect training score has to be memorization:

import numpy as np
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split

X, y = make_classification(
    n_samples=2000, n_features=12, n_informative=4, n_redundant=2,
    flip_y=0.1, random_state=0,
)
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.3, stratify=y, random_state=0
)

full = DecisionTreeClassifier(random_state=0).fit(X_train, y_train)
print("depth:", full.get_depth(), "| leaves:", full.get_n_leaves())
print("train accuracy:", round(full.score(X_train, y_train), 3))
print("test accuracy: ", round(full.score(X_test, y_test), 3))

With no limits, the tree grows until every training sample sits in a pure leaf. Expect a training accuracy of 1.000, a test accuracy well below it, and a tree with hundreds of leaves, many of them fitted to noise. That gap is the textbook picture of overfitting, and pruning exists to close it.

Pre-Pruning: Stop the Tree Early

Pre-pruning (early stopping) sets limits that prevent the tree from growing too far. The main controls:

• max_depth: the longest allowed path from root to leaf.
• min_samples_split: the minimum samples a node needs to be considered for splitting.
• min_samples_leaf: the minimum samples allowed in a leaf. This is often the most effective single knob, because it stops tiny, noisy leaves.
• max_leaf_nodes: a cap on total leaves.
• min_impurity_decrease: refuses splits that don't improve purity enough.

Rather than guessing, search for good values with cross-validation on the training set:

from sklearn.model_selection import GridSearchCV

param_grid = {
    "max_depth": [2, 3, 4, 5, 6, 8, None],
    "min_samples_leaf": [1, 5, 10, 20, 50],
}
search = GridSearchCV(
    DecisionTreeClassifier(random_state=0), param_grid, cv=5, scoring="accuracy"
)
search.fit(X_train, y_train)

print("best params:", search.best_params_)
print("cv accuracy:", round(search.best_score_, 3))
print("test accuracy:", round(search.score(X_test, y_test), 3))

Pre-pruning is simple and fast. Its weakness is that it can stop too soon: a split that looks useless on its own might enable a very useful split below it, and a greedy stopping rule never finds out.

Post-Pruning: Grow Big, Then Cut Back

Post-pruning takes the opposite approach: grow a large tree, then remove branches that don't earn their keep. Scikit-learn implements minimal cost-complexity pruning, which minimizes:

R_α(T) = R(T) + α · |T|

Here R(T) is the tree's total training error, |T| is the number of leaves, and α is a penalty per leaf. At α = 0 you keep the full tree, and as α grows, the tree is cut back more aggressively, down to a single node. The ccp_alpha parameter sets α, and the method cost_complexity_pruning_path computes the exact values of α at which the tree changes shape.

Choose α by cross-validation:

from sklearn.model_selection import cross_val_score

path = DecisionTreeClassifier(random_state=0).cost_complexity_pruning_path(X_train, y_train)
alphas = np.unique(path.ccp_alphas)[:-1]            # drop the last alpha: it prunes to one node
alphas = alphas[:: max(1, len(alphas) // 40)]        # thin the grid to keep this fast

cv_scores = [
    cross_val_score(DecisionTreeClassifier(random_state=0, ccp_alpha=a),
                    X_train, y_train, cv=5)
    for a in alphas
]
means = np.array([s.mean() for s in cv_scores])
ses = np.array([s.std(ddof=1) / np.sqrt(len(s)) for s in cv_scores])

best = means.argmax()
one_se_alpha = alphas[np.where(means >= means[best] - ses[best])[0].max()]

print("alpha with best CV score:", round(alphas[best], 5))
print("one-standard-error alpha:", round(one_se_alpha, 5))

The last line applies the one-standard-error rule: instead of the alpha with the single highest CV score, take the largest alpha whose score is within one standard error of the best. That picks a simpler tree with statistically indistinguishable accuracy, which tends to generalize better and is easier to read. Now fit and compare:

pruned = DecisionTreeClassifier(random_state=0, ccp_alpha=one_se_alpha).fit(X_train, y_train)

for name, model in [("full", full), ("pruned", pruned)]:
    print(f"{name:>6}: depth={model.get_depth():>2}  leaves={model.get_n_leaves():>3}  "
          f"train={model.score(X_train, y_train):.3f}  test={model.score(X_test, y_test):.3f}")

You should see the pruned tree use a small fraction of the leaves, with a lower training score, a higher or equal test score, and a much smaller gap between the two. You can also plot mean CV accuracy against alpha to see the tradeoff:

plt.plot(alphas, means, marker="o")
plt.xscale("symlog", linthresh=1e-4)
plt.xlabel("ccp_alpha"); plt.ylabel("mean CV accuracy")
plt.axvline(one_se_alpha, linestyle="--")
plt.show()

Pre-Pruning vs Post-Pruning

Pre-pruning is cheaper and easy to explain, while post-pruning is more principled because it considers the whole grown tree before deciding what to cut. Many practitioners combine them: cap min_samples_leaf to avoid pointless tiny leaves, then tune ccp_alpha. Whichever you use, select the settings with cross-validation on training data and keep the test set for a single final check — as covered in the cross-validation guide in this series.

Feature Importance, and Its Catch

A fitted tree exposes feature_importances_, the total impurity reduction attributable to each feature. It's convenient, but impurity-based importance is biased toward continuous or high-cardinality features and is computed on training data, so it can credit features that only helped memorize noise. A more trustworthy view comes from permutation importance on held-out data:

from sklearn.inspection import permutation_importance

result = permutation_importance(pruned, X_test, y_test, n_repeats=20, random_state=0)
for i in result.importances_mean.argsort()[::-1][:5]:
    print(f"feature {i}: {result.importances_mean[i]:.3f} ± {result.importances_std[i]:.3f}")

Permutation importance shuffles one feature at a time and measures how much test performance drops. As with every explanation method covered in this series — from SHAP values to partial dependence — it describes what the model uses, not what causes the outcome.

Regression Trees and Other Knobs

DecisionTreeRegressor works the same way, but splits minimize squared error and each leaf predicts the average target of its samples. Pruning applies identically through ccp_alpha and the pre-pruning parameters. Regression trees produce piecewise-constant predictions, so they can't extrapolate beyond the range of training targets.

A few other parameters are worth knowing:

• class_weight="balanced" reweights classes for imbalanced data, which shifts the tree toward the minority class.
• max_features restricts the features considered at each split. It's the randomization ingredient that random forests rely on.
• Recent scikit-learn versions also support native handling of missing values in decision trees and monotonic constraints on features. Check the documentation for your installed version, since support depends on the release.

Strengths and Weaknesses

Strengths: easy to understand and explain, no feature scaling needed, handles numeric and (encoded) categorical features, captures non-linear relationships and interactions, and fast to train and predict.

Weaknesses: high variance, since a small change in the data can produce a very different tree. Splits are axis-aligned, so diagonal boundaries need many splits. They can't extrapolate. Impurity-based importance is biased. And a deep tree stops being interpretable, since a 40-level tree is no more readable than a black box.

That variance problem is exactly why trees are rarely the final model on tabular data. Combine many of them and you get random forests, which average out the variance, or gradient boosting, which fits trees sequentially to correct earlier errors, covered in the algorithms guide in this series. A single pruned tree earns its place when you need a small, auditable rule set, or as a baseline — much like logistic regression serves as a strong first classifier.

Putting a Tree Into Production

If you need to explain or reimplement the model, export_text gives you the rules, and a pruned tree with a handful of leaves can be turned into plain if/else logic. For serving, persist the fitted model with joblib as with any scikit-learn estimator. Version the pruning parameters alongside the model, and re-validate when data drifts, because tree structure can change dramatically when retrained on new data.

Common Pitfalls

• Evaluating an unpruned tree on its own training data and calling it accurate.
• Tuning pruning parameters on the test set instead of with cross-validation.
• Reading feature_importances_ from an overfit tree as ground truth.
• Using accuracy on imbalanced data without checking precision and recall.
• Assuming a deeper tree is a better tree — depth buys fit, not generalization.
• Presenting a very deep tree as "interpretable" when nobody can read it.
• Expecting a regression tree to extrapolate beyond the target range it saw.
• Forgetting that two nearly identical datasets can yield very different trees.

Unlock AI That Actually Works

Get lifetime access to GPT-6 Astra, Claude Fable 5.1, Gemini 3.5, Grok 4.5, and more — all in one platform. Build websites, apps, videos, content, and digital products from a single command. No monthly fees. No tool-hopping.

Click here to get GPTAstra Max now — one-time payment, lifetime access.

Frequently Asked Questions

What is the difference between pre-pruning and post-pruning?

Pre-pruning stops the tree while it grows by setting limits such as max_depth, min_samples_leaf, and max_leaf_nodes — cheap and simple, but it can stop too early, because a split that looks useless alone might enable a useful split below it. Post-pruning grows the full tree first, then removes branches whose contribution to accuracy does not justify their leaf count, which is more principled because it judges the whole tree before cutting. Many practitioners combine both: cap min_samples_leaf to avoid pointless tiny leaves, then tune ccp_alpha.

How do I choose the ccp_alpha value?

Let cross-validation pick it. cost_complexity_pruning_path returns the exact alpha values at which the tree changes shape; drop the final alpha (it prunes to a single node), evaluate the rest with cross-val on the training set, and take either the best or — recommended — the largest alpha within one standard error of the best score. That one-standard-error rule favors the simpler tree with statistically indistinguishable accuracy, which tends to generalize better and is easier to read.

Why does my decision tree show 100% training accuracy?

Because an unconstrained tree keeps splitting until every training sample sits in a pure leaf, so it always fits the training data perfectly — including any label noise. On the noisy synthetic dataset in this article you get a train accuracy of 1.000, a test accuracy well below it, and hundreds of leaves fitted to noise. That gap is overfitting; fix it with pre-pruning limits or cost-complexity pruning selected by cross-validation.

Which feature importance should I trust?

Prefer permutation importance measured on held-out data over the built-in feature_importances_. Impurity-based importance is biased toward continuous or high-cardinality features and is computed on training data, so it can credit features that only helped memorize noise. Permutation importance shuffles one feature at a time and measures how much test performance drops, which reflects what the model actually uses — though, like every explanation method, it describes the model, not causes in the world.

Wrapping This Up

A decision tree learns a flowchart by greedily choosing splits that reduce impurity. Left unconstrained it memorizes noise, so you control its size with pre-pruning parameters such as max_depth and min_samples_leaf, or with post-pruning through cost-complexity pruning and ccp_alpha, choosing values by cross-validation and favoring the simpler tree within one standard error of the best. Visualize with plot_tree and export_text, and judge feature importance with permutation importance on held-out data.

Should a single tree be your final model? Usually not for accuracy, because ensembles beat it on most tabular problems. But for a readable baseline, a rule set you can audit, or a first look at which features matter, a well-pruned tree is hard to beat. Try the code above, plot the accuracy against alpha, and watch what happens as the tree shrinks.

Related Articles

• Machine Learning Algorithms Explained: Complete Beginner's Guide
• Logistic Regression Explained: Build a Classifier Step by Step
• SHAP Values in Python: Explain Your Model's Predictions
• Scikit-learn Cross-Validation: Master K-Fold Testing