Decision Trees
Machine Learning
Algorithm Implementation
Data Science
Feature Selection

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

  1. 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 KK classes, Gini impurity is G=1k=1Kpk2G = 1 - \sum_{k=1}^{K} p_k^2, where pkp_k is the proportion of class kk.
    • Entropy: Used in the context of information gain, it quantifies the amount of uncertainty or disorder in a dataset. Entropy is H=k=1Kpklog2(pk)H = -\sum_{k=1}^{K} p_k \log_2(p_k).
    • Variance Reduction: In regression tasks, the focus is typically on reducing the variance within each subset.
  2. 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.
  3. 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.
  4. 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

  1. Start with the Root Node:
    • For each candidate feature, calculate the impurity of the root node (entropy, Gini, or variance).
  2. 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.
  3. 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):
    IG(D,A)=I(D)vvalues(A)DvDI(Dv)IG(D, A) = I(D) - \sum_{v \in values(A)} \frac{|D_v|}{|D|} I(D_v)
    • Where IG(D,A)IG(D, A) is the information gain for dataset DD using feature AA, I(D)I(D) is the impurity of dataset DD, and DvD_v represents the subset of DD with feature AA having value vv.
  4. 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.
  5. 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.

WeatherTemperatureHumidityWindPlay Tennis
SunnyHotHighWeakNo
SunnyHotHighStrongNo
OvercastHotHighWeakYes
RainMildHighWeakYes
RainCoolNormalWeakYes

First, calculate the entropy of the target variable (Play Tennis). With 2 "No" and 3 "Yes" outcomes:

H(D)=25log22535log2350.971H(D) = -\frac{2}{5}\log_2\frac{2}{5} - \frac{3}{5}\log_2\frac{3}{5} \approx 0.971

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 H(D)H(D).

Summary of Feature Gains

FeatureInformation GainSelected (Yes/No)
Weather0.246Yes
Temperature0.029No
Humidity0.151No
Wind0.048No

"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: IG(D,A)=I(D)vDvDI(Dv)IG(D, A) = I(D) - \sum_{v} \frac{|D_v|}{|D|} I(D_v). 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.


Course illustration
Course illustration

All Rights Reserved.