Decision Tree Algorithm in Machine Learning: A Beginner-Friendly Guide with Python

Introduction

Machine Learning is a branch of Artificial Intelligence that allows computers to learn patterns from data and make predictions without being explicitly programmed for every task.

One of the most popular supervised machine learning algorithms is the Decision Tree Algorithm.

A Decision Tree is a powerful algorithm used for both classification and regression problems. It makes decisions by splitting a dataset into smaller groups based on different conditions.

It works similarly to how humans make decisions in everyday life.

For example, when deciding whether to play cricket, we might consider:

• Is the weather sunny?
• Is it raining?
• Is the ground wet?
• Is the temperature suitable?

Based on these conditions, we make a final decision.

In this article, we will understand what a Decision Tree is, how it works, its mathematical concepts, different types, advantages, disadvantages, and how to implement it using Python.

1. What is a Decision Tree?

A Decision Tree is a supervised machine learning algorithm used for classification and regression tasks.

It represents decisions in a tree-like structure, where each internal node represents a condition, each branch represents an outcome, and each leaf node represents the final prediction.

A Decision Tree consists of three important components:

1. Root Node

The root node represents the starting point of the decision tree. It contains the entire dataset before splitting.

2. Decision Node

A decision node represents a condition or test performed on a feature to divide the dataset into smaller groups.

3. Leaf Node

A leaf node represents the final output or prediction of the algorithm.

Example

Suppose we want to predict whether a person will play cricket based on the weather.

WeatherPlay CricketSunnyNoRainyNoCloudyYesSunnyNoCloudyYes

The Decision Tree learns patterns from this dataset and creates conditions to predict whether someone will play cricket.

2. How Does a Decision Tree Work?

A Decision Tree works by repeatedly splitting the dataset into smaller subsets based on the best features.

The algorithm follows these steps:

1. Start with the complete training dataset.
2. Identify the available features.
3. Calculate the best feature for splitting.
4. Divide the dataset based on that feature.
5. Repeat the splitting process for each subset.
6. Continue until a stopping condition is reached.
7. Generate the final prediction at the leaf nodes.

The main objective is to create groups that contain similar data points.

For example, consider a dataset used to predict whether a customer will purchase a product.

The algorithm may create conditions such as:

• Is the customer’s age greater than 25?
• Is the customer’s income greater than ₹30,000?
• Has the customer purchased before?

Based on these conditions, the algorithm predicts whether the customer will purchase the product.

3. Important Concepts in Decision Trees

Decision Trees use different mathematical techniques to identify the best feature for splitting.

The most important concepts are:

A. Entropy

Entropy measures the impurity or randomness in a dataset.

If all records belong to the same class, entropy is zero.

If the classes are equally distributed, entropy is higher.

The formula is:

Entropy(S) = -Σ p(i) log₂ p(i)

Where:

• S = Dataset
• p(i) = Probability of each class
• log₂ = Logarithm with base 2

Example:

Suppose a dataset contains:

• 5 Yes
• 5 No

Probability of Yes = 0.5

Probability of No = 0.5

Entropy = -(0.5 log₂ 0.5 + 0.5 log₂ 0.5)

Entropy = 1

This indicates maximum impurity for a binary classification problem.

B. Information Gain

Information Gain measures how much uncertainty is reduced after splitting a dataset using a particular feature.

The algorithm prefers the feature that provides the highest Information Gain.

Formula:

Information Gain = Entropy(Parent) — Weighted Entropy(Children)

Where:

• Parent Entropy = Entropy before splitting.
• Children Entropy = Weighted average entropy after splitting.

Example:

Suppose the parent entropy is 1.

After splitting, the weighted child entropy is 0.4.

Information Gain = 1–0.4

Information Gain = 0.6

Therefore, the split reduces uncertainty by 0.6.

C. Gini Impurity

Gini Impurity is another method used to measure the impurity of a dataset.

It calculates the probability of incorrectly classifying a randomly selected data point if it is labelled according to the class distribution in the dataset.

Formula:

Gini = 1 — Σ p(i)²

Where:

• p(i) = Probability of each class.

Example:

Suppose a dataset contains:

• 5 Yes
• 5 No

Gini = 1 — (0.⁵² + 0.⁵²)

Gini = 1 — (0.25 + 0.25)

Gini = 0.5

A lower Gini value indicates a purer dataset.

D. Splitting Criteria

Decision Trees use splitting criteria to identify the best feature for dividing the dataset.

Common splitting criteria include:

• Entropy
• Information Gain
• Gini Impurity

In Scikit-learn, the DecisionTreeClassifier supports criteria such as:

• gini
• entropy
• log_loss

4. Types of Decision Trees

There are two major types of Decision Trees based on their prediction tasks.

A. Classification Tree

A Classification Tree is used when the target variable contains categorical values.

Examples:

• Spam or Not Spam
• Pass or Fail
• Yes or No
• Customer will buy or not buy

Python class:

DecisionTreeClassifier()

B. Regression Tree

A Regression Tree is used when the target variable contains continuous numerical values.

Examples:

• House price prediction
• Salary prediction
• Temperature prediction
• Sales prediction

Python class:

Get Mohithsrinath’s stories in your inbox

Join Medium for free to get updates from this writer.

Remember me for faster sign in

DecisionTreeRegressor()

5. Real-World Applications of Decision Trees

Decision Trees are widely used in different industries.

1. Healthcare

Predicting diseases based on patient information and symptoms.

2. Banking

Predicting whether a customer is eligible for a loan.

3. E-commerce

Predicting customer purchasing behaviour.

4. Fraud Detection

Identifying potentially fraudulent transactions based on transaction characteristics.

5. Education

Predicting student performance based on attendance, marks, and other factors.

6. Customer Segmentation

Grouping customers based on their characteristics and behaviour.

6. Implementing Decision Tree Using Python

Let’s implement a Decision Tree Classification model using the famous Iris dataset.

We will use Python and Scikit-learn.

Step 1: Import Required Libraries

import pandas as pd
import numpy as np

from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier
from sklearn.metrics import accuracy_score
from sklearn.metrics import classification_report
from sklearn.metrics import confusion_matrix

Step 2: Load the Dataset

iris = load_iris()

X = iris.data
y = iris.target
print(X[:5])
print(y[:5])

Explanation:

• X represents the independent features.
• y represents the dependent or target variable.
• The Iris dataset contains four numerical features.
• The target contains three flower species.

Step 3: Split the Dataset

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

Explanation:

• 80% of the dataset is used for training.
• 20% of the dataset is used for testing.
• random_state=42 ensures reproducibility.

Step 4: Create the Decision Tree Model

model = DecisionTreeClassifier(
    criterion="gini",
    random_state=42
)

Explanation:

• DecisionTreeClassifier creates the classification model.
• criterion=”gini” uses Gini Impurity to select splits.
• random_state=42 makes the model reproducible.

Step 5: Train the Model

model.fit(X_train, y_train)

The model learns patterns from the training dataset and identifies the best splitting conditions.

Step 6: Make Predictions

y_pred = model.predict(X_test)

print(y_pred)

The model predicts the flower species using the testing dataset.

Step 7: Evaluate Model Performance

accuracy = accuracy_score(y_test, y_pred)

print("Accuracy:", accuracy)

We can also generate a classification report:

print(classification_report(y_test, y_pred))

And a confusion matrix:

print(confusion_matrix(y_test, y_pred))

These metrics help us understand how accurately the model predicts the flower species.

7. Visualizing a Decision Tree

One of the biggest advantages of a Decision Tree is that we can visualize its decision-making process.

Scikit-learn provides a function called plot_tree().

Example:

import matplotlib.pyplot as plt
from sklearn.tree import plot_tree

plt.figure(figsize=(15, 10))
plot_tree(
    model,
    feature_names=iris.feature_names,
    class_names=iris.target_names,
    filled=True
)
plt.show()

Explanation:

• plt.figure() defines the figure size.
• plot_tree() visualizes the trained decision tree.
• feature_names displays the feature names.
• class_names displays the target classes.
• filled=True adds colours to the tree nodes.

This visualization helps us understand how the algorithm makes predictions.

8. Overfitting in Decision Trees

Overfitting is one of the major problems in Decision Tree algorithms.

It happens when a model learns the training data too well, including noise and unnecessary patterns.

As a result:

• Training accuracy becomes very high.
• Testing accuracy may become low.
• The model may perform poorly on unseen data.

How to Reduce Overfitting?

We can reduce overfitting using different techniques:

1. max_depth

Limits the maximum depth of the tree.

model = DecisionTreeClassifier(max_depth=3)

2. min_samples_split

Defines the minimum number of samples required to split an internal node.

model = DecisionTreeClassifier(min_samples_split=5)

3. min_samples_leaf

Defines the minimum number of samples required at a leaf node.

model = DecisionTreeClassifier(min_samples_leaf=3)

4. Pruning

Pruning removes unnecessary branches from a decision tree to make the model simpler and improve generalization.

Example:

model = DecisionTreeClassifier(ccp_alpha=0.01)

The ccp_alpha parameter controls cost-complexity pruning. A larger value generally produces a smaller tree.

9. Advantages of Decision Tree

1. Easy to understand and interpret.
2. Can be used for classification and regression.
3. Handles numerical and categorical features.
4. Requires minimal data preprocessing.
5. Does not require feature scaling or normalization.
6. Can capture non-linear relationships.
7. Provides visual representations of decisions.

10. Disadvantages of Decision Tree

1. Can easily overfit the training data.
2. Small changes in data can produce a different tree.
3. Deep trees can become difficult to interpret.
4. A greedy splitting strategy may not find the globally optimal tree.
5. A single Decision Tree may not provide the best prediction performance compared with ensemble methods.
6. Imbalanced datasets can affect model performance.

11. Decision Tree vs Naive Bayes

Both algorithms are popular supervised machine learning algorithms, but they work differently.

FeatureDecision TreeNaive BayesAlgorithm typeTree-basedProbability-basedMain conceptDecision rulesBayes’ TheoremPredictionSplitting conditionsClass probabilitiesFeature independenceNot requiredAssumedOverfittingCan overfitUsually less prone to memorizing training dataCommon applicationLoan predictionSpam detectionInterpretabilityEasy to visualizeEasy to understand probabilistically

Both algorithms can solve classification problems, but the best choice depends on the dataset and the problem.

12. Conclusion

The Decision Tree Algorithm is one of the most important supervised machine learning algorithms.

It makes predictions by dividing data into smaller groups using different conditions and decision rules.

It is useful for both classification and regression problems and is widely applied in healthcare, banking, e-commerce, fraud detection, and many other industries.

Understanding concepts such as Entropy, Information Gain, Gini Impurity, and Pruning helps us understand how Decision Trees select the best features and make predictions.

If you are beginning your Machine Learning journey, learning and implementing Decision Trees using Python and Scikit-learn is an excellent way to build a strong foundation.

Keep learning, keep practicing, and keep building!

Tags

Machine Learning, Decision Tree, Python, Data Science, Artificial Intelligence