Decision tree implementation for returning the next feature to split the tree
Master System Design with Codemia
Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.
Introduction
In machine learning, decision trees are one of the most intuitive and interpretable models for classification and regression tasks. One critical aspect of building a decision tree is determining the next feature to split the data at each node. This decision influences the tree's performance, complexity, and interpretability.
Understanding Decision Trees
A decision tree is a hierarchical structure where internal nodes represent tests on features, branches represent the outcomes of these tests, and leaf nodes represent the final prediction. The objective is to create a tree that most effectively separates the data into distinct classes or predicts a target variable, achieving high accuracy with minimal complexity.
Selecting the Next Feature to Split
Selecting the next feature to split involves determining which feature, when used for splitting, results in the most "pure" or distinct subsequent sets of data. Purity refers to the homogeneity of a dataset's classes in classification tasks or the reduction in variance for regression tasks.
Key Concepts
- Impurity Measures:
- Gini Impurity: Measures the likelihood of an incorrect classification if a random instance from the dataset is labeled according to the distribution of labels in the subset. For a node with classes, Gini impurity is , where is the proportion of class .
- Entropy: Used in the context of information gain, it quantifies the amount of uncertainty or disorder in a dataset. Entropy is .
- Variance Reduction: In regression tasks, the focus is typically on reducing the variance within each subset.
- Information Gain:
- Information gain is calculated using entropy and represents the decrease in entropy (or impurity) achieved by partitioning the data on a specific feature.
- Gain Ratio:
- A normalization of information gain, it accounts for the intrinsic information of a split, balancing feature bias with the actual reduction in impurity. The C4.5 algorithm uses gain ratio to avoid favoring features with many distinct values.
- Reduction in Variance:
- For regression trees, choosing the feature that leads to the largest reduction in variance from the parent node to child nodes is the standard criterion (used by CART for regression).
Algorithm for Selecting the Next Feature
- Start with the Root Node:
- For each candidate feature, calculate the impurity of the root node (entropy, Gini, or variance).
- Calculate Impurity After Splitting:
- For each feature, compute the impurity after splitting the data based on possible values of the feature. This results in child nodes representing sub-datasets.
- Calculate Information Gain:
- Information gain, for a feature, is the difference between the impurity of the original dataset and the weighted sum of impurities for each subset (child node):
- Where is the information gain for dataset using feature , is the impurity of dataset , and represents the subset of with feature having value .
- Choose the Feature with the Highest Gain:
- The feature with the maximum information gain (or highest gain ratio for C4.5) is chosen as the splitting feature.
- Repeat Recursively:
- This process is repeated for each child node until a stopping criterion is met (for example, maximum depth, minimum samples per leaf, or no further information gain).
Example
Consider a simple dataset with features such as weather conditions and a target to predict whether a person will play tennis.
| Weather | Temperature | Humidity | Wind | Play Tennis |
| Sunny | Hot | High | Weak | No |
| Sunny | Hot | High | Strong | No |
| Overcast | Hot | High | Weak | Yes |
| Rain | Mild | High | Weak | Yes |
| Rain | Cool | Normal | Weak | Yes |
First, calculate the entropy of the target variable (Play Tennis). With 2 "No" and 3 "Yes" outcomes:
Then compute information gain for each feature. For example, splitting on "Weather" creates three subsets (Sunny, Overcast, Rain), and the weighted entropy of those subsets is subtracted from .
Summary of Feature Gains
| Feature | Information Gain | Selected (Yes/No) |
| Weather | 0.246 | Yes |
| Temperature | 0.029 | No |
| Humidity | 0.151 | No |
| Wind | 0.048 | No |
"Weather" has the highest information gain and is selected for the first split.
Additional Details
Handling Continuous Features
To handle continuous features, data is split based on a threshold value. The algorithm evaluates all possible thresholds (typically midpoints between sorted adjacent values) and selects the one producing the highest information gain. This converts a continuous feature into a binary split.
Pruning
To prevent overfitting, trees are often pruned after construction. Pre-pruning sets stopping criteria during tree building (max depth, min samples). Post-pruning builds the full tree first, then removes branches that do not improve validation accuracy.
Advantages and Disadvantages
- Advantages: Easy to interpret and visualize, require little data preprocessing, can handle both numerical and categorical data, and capture non-linear relationships.
- Disadvantages: Prone to overfitting (mitigated by pruning or ensembles), small changes in data can lead to significantly different trees, and they tend to be biased toward features with more levels when using information gain alone.
Summary
The process of selecting the next feature to split a decision tree is crucial for building an effective model. The key formula is information gain: . Through understanding and implementing concepts like information gain, gain ratio, and variance reduction, you can construct decision trees that balance accuracy and interpretability. Modern implementations in libraries like scikit-learn handle these calculations automatically, but understanding the underlying mechanics helps you tune hyperparameters and debug model behavior effectively.

