Decision trees are one of the most intuitive algorithms in artificial intelligence. They make predictions by asking a sequence of questions: is this value above a threshold, does this category match, which branch should this observation follow next?
That structure makes decision trees useful for learning machine learning because they are visual, explainable, and close to how humans already reason through choices. A decision tree can classify an email as spam or not spam, predict whether a customer might churn, estimate a house price, or help explain which variables drive a model’s output.
What is a decision tree in AI?
A decision tree is a supervised machine learning model that splits data into smaller and smaller groups until it can make a prediction. Each internal node represents a question about a feature, each branch represents an answer, and each leaf represents the final prediction.
For example, a simple decision tree for predicting whether someone might buy a bike could ask:
- Is the person younger than 35?
- Do they live within 5 km of work?
- Have they bought sports equipment before?
The model learns these questions from data. It does not need a human to write every rule manually. That is what makes it machine learning rather than a traditional hand-coded decision table.
Decision trees: a simple yet powerful tool
Picture a tree with branches that represent choices and leaves that represent conclusions. This visual representation makes the decision-making process easy to follow. Starting from the root, the tree evaluates the data and moves down the branches, making decisions based on feature values.
This intuitive algorithm is a jack-of-all-trades. It can handle both classification and regression tasks, and it often works well as a first model when you want to understand a dataset.
How Does It Work?
Decision trees work by splitting data into distinct groups based on specific features. Each split creates new branches, leading to more specific decisions. This process is like a thoughtful interview, asking precise questions to reach a useful conclusion.
During training, the algorithm searches for splits that make the resulting groups more homogeneous. In a classification problem, that means each leaf should contain mostly one class. In a regression problem, that means each leaf should contain values that are close to each other.
Common splitting criteria include:
- Gini impurity: measures how mixed the classes are in a node.
- Entropy: measures uncertainty in the class distribution.
- Information gain: measures how much a split reduces uncertainty.
- Mean squared error: often used for regression trees.
The tree keeps splitting until it reaches a stopping rule, such as a maximum depth, a minimum number of samples in a leaf, or no useful improvement from additional splits.
Classification trees vs regression trees
Decision trees can solve two common types of supervised learning problems.
| Type | Goal | Example |
|---|---|---|
| Classification tree | Predict a category | Will a customer churn: yes or no? |
| Regression tree | Predict a number | What will the house price be? |
A classification tree returns a class label or class probability. A regression tree returns a numerical prediction, usually based on the average target value in the final leaf.
The Versatility of Decision Trees
The best part about decision trees? Their versatility knows no bounds. They handle numerical features naturally, and with the right preprocessing they can handle categorical variables too. They are also easy to visualise, giving us a clear understanding of the decision-making process. It’s like having a colleague who can explain complex decisions in a simple, step-by-step manner.
Decision trees also form the foundation for more powerful algorithms such as random forests, gradient boosted trees, XGBoost, LightGBM, and CatBoost. If you understand a single decision tree, those ensemble methods become much easier to learn.
Strengths and weaknesses
Decision trees are popular because they are:
- easy to explain
- able to model non-linear relationships
- useful for both classification and regression
- less sensitive to feature scaling than many other algorithms
- good at surfacing important variables
But they also have weaknesses:
- they can overfit if allowed to grow too deep
- small changes in the data can produce a different tree
- a single tree may be less accurate than ensemble methods
- they can create biased splits when features have many possible values
This is why a decision tree is often a great starting point, while a random forest or gradient boosted model is often a stronger production choice.
Overfitting and pruning
A decision tree can keep splitting until it memorises the training data. That looks impressive during training, but it usually performs badly on new data.
To avoid this, control the complexity of the tree. In scikit-learn, useful parameters include:
max_depth: limits how deep the tree can growmin_samples_split: requires enough samples before a node can splitmin_samples_leaf: requires enough samples in each final leafccp_alpha: applies cost-complexity pruning
These parameters make the tree simpler and usually improve how well it generalises.
Examples in Action
Let’s translate these concepts into a practical example using Python. This script trains a decision tree classifier on the Iris dataset with scikit-learn.
from sklearn.tree import DecisionTreeClassifier
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
iris = load_iris()
X_train, X_test, y_train, y_test = train_test_split(
iris.data,
iris.target,
test_size=0.2,
random_state=0,
)
tree = DecisionTreeClassifier(
criterion="gini",
max_depth=3,
random_state=0,
)
tree.fit(X_train, y_train)
predictions = tree.predict(X_test)
accuracy = accuracy_score(y_test, predictions)
print(predictions)
print(f"Accuracy: {accuracy:.2f}")This code uses a maximum depth of 3 so the tree stays small enough to interpret. That is often a good first move: start with a simple tree, understand the decisions it makes, then increase complexity only if the results justify it.
Prefer to run the example yourself? Open the companion notebook in Notebook Studio or download it locally.
When should you use a decision tree?
Use a decision tree when:
- you need an interpretable model
- you want a quick baseline
- feature relationships are likely non-linear
- stakeholders need a visual explanation
- you are preparing to learn random forests or gradient boosting
Avoid relying on a single tree when prediction performance is the only goal and interpretability is less important. In that case, tree ensembles often perform better.
Key Takeaways
Decision trees are a robust and versatile tool for anyone tackling machine learning projects. Their ability to split complex datasets into understandable decisions makes them useful for both learning and applied data science.
The most important things to remember:
- A decision tree predicts by following a path of learned questions.
- Classification trees predict categories; regression trees predict numbers.
- Splitting criteria such as Gini impurity and entropy help the tree choose useful questions.
- Tree depth and pruning matter because unrestricted trees can overfit.
- Decision trees are the building blocks behind random forests and gradient boosted trees.
Recommendations
If you like this article, you might also find this one interesting
Peak Performance Modeling: Harnessing Bagging and Boosting for Superior Results
In the world of machine learning, two key ensemble techniques stand out. They can improve model performance. They are: bagging and boosting. But how do you decide which one to use? Let’s break it down in my article on bagging and boosting for better model performance.
References
Let’s Connect
I’d love to hear your thoughts on this fascinating topic!
Feel free to connect with me on Twitter at https://twitter.com/feddernico or explore more of my content on Medium https://medium.com/@federico.viscioletti and Substack https://feddernico.substack.com
Happy learning!