Decision Trees From Scratch in Python: The Math of Gini Impurity and Splits

A Decision Tree is a non-parametric supervised learning algorithm that partitions data by recursively asking a sequence of binary yes/no questions about feature values, branching at each internal node until arriving at a terminal leaf prediction.

Think about playing the game 'Twenty Questions' to identify an animal: 'Is it a mammal?' Yes. 'Does it live in water?' Yes. 'Does it have a blowhole?' Yes. Within three questions, you narrow the candidates from millions of possibilities straight to a whale. A decision tree does the exact same thing with datasets: it searches through all features and numerical thresholds to find the questions that split the data into the purest possible subsets.

1. Why Decision Trees Are Fundamental

• Human Interpretability: Decision trees mirror human diagnostic rules. You can visualize and trace the exact logic path behind every single prediction.
• Zero Feature Scaling Required: Because splits compare values strictly using inequalities ($x_j \le \text{threshold}$), monotonicity is preserved regardless of numerical feature scale.
• Non-Linear Decision Boundaries: Captures complex interactions and step-like boundaries that linear models cannot separate.
• Foundational Building Blocks: Advanced ensemble models like Random Forests, Extra Trees, and XGBoost are built from large collections of decision trees.

2. The Splitting Criteria Decision Matrix

3. The Math of How Splits Are Chosen

1. Gini Impurity Formulation

Gini Impurity measures the probability of misclassifying a randomly chosen element from the set if it were randomly labeled according to the class distribution:

$$\text{Gini} = 1 - \sum_{k=1}^{K} p_k^2$$

• Pure Node ($p_1 = 1.0, p_2 = 0.0$): $\text{Gini} = 1 - (1.0^2 + 0.0^2) = 0.0$ (Zero impurity).
• Maximally Impure ($p_1 = 0.5, p_2 = 0.5$): $\text{Gini} = 1 - (0.5^2 + 0.5^2) = 0.50$ (Maximum chaos for binary classification).

2. Information Gain (IG)

Information Gain measures the reduction in impurity achieved by splitting parent node $D$ into left child $D_L$ and right child $D_R$:

$$\text{IG} = \text{Gini}_{\text{parent}} - \left( \frac{m_L}{m} \text{Gini}_L + \frac{m_R}{m} \text{Gini}_R \right)$$

3. Regression Trees: Variance Reduction

For continuous target values, the tree minimizes Mean Squared Error (variance) rather than class impurity:

$$\Delta \text{MSE} = \text{MSE}_{\text{parent}} - \left( \frac{m_L}{m} \text{MSE}_L + \frac{m_R}{m} \text{MSE}_R \right)$$

In a regression tree leaf, the prediction $\hat{y}$ is simply the mean value $\bar{y}$ of all training samples assigned to that leaf.

4. Worked Numerical Example (By Hand)

Suppose we have $m = 10$ marbles (6 Red, 4 Blue). We evaluate a candidate split size > 5mm:

• Left Child ($m_L = 5$): 4 Red, 1 Blue $\implies p_R = 0.8, p_B = 0.2$.
• Right Child ($m_R = 5$): 2 Red, 3 Blue $\implies p_R = 0.4, p_B = 0.6$.
• Parent Gini: $\text{Gini}_{\text{parent}} = 1 - (0.6^2 + 0.4^2) = 1 - 0.52 = 0.48$.
• Left Gini: $\text{Gini}_L = 1 - (0.8^2 + 0.2^2) = 1 - 0.68 = 0.32$.
• Right Gini: $\text{Gini}_R = 1 - (0.4^2 + 0.6^2) = 1 - 0.52 = 0.48$.
• Weighted Child Impurity: $\frac{5}{10}(0.32) + \frac{5}{10}(0.48) = 0.16 + 0.24 = 0.40$.
• Information Gain: $\text{IG} = 0.48 - 0.40 = 0.08$.

The algorithm computes this exact score for every unique feature threshold and selects the one with the highest $\text{IG}$.

5. Code: Decision Tree From Scratch in Python

1. NumPy Tree Building Engine From Scratch

import numpy as np

def calculate_gini(y):
    if len(y) == 0:
        return 0.0
    proportions = np.array([np.mean(y == c) for c in np.unique(y)])
    return 1.0 - np.sum(proportions ** 2)

def calculate_information_gain(y, y_left, y_right):
    m = len(y)
    weighted_child = (len(y_left) / m) * calculate_gini(y_left) + (len(y_right) / m) * calculate_gini(y_right)
    return calculate_gini(y) - weighted_child

def find_best_split(X, y):
    best_gain = -1.0
    best_feature, best_threshold = None, None
    m, n_features = X.shape

    for feature_idx in range(n_features):
        thresholds = np.unique(X[:, feature_idx])
        for t in thresholds:
            left_mask = X[:, feature_idx] <= t
            right_mask = ~left_mask

            if left_mask.sum() == 0 or right_mask.sum() == 0:
                continue

            gain = calculate_information_gain(y, y[left_mask], y[right_mask])
            if gain > best_gain:
                best_gain = gain
                best_feature = feature_idx
                best_threshold = t

    return best_feature, best_threshold, best_gain

class TreeNode:
    def __init__(self, feature=None, threshold=None, left=None, right=None, prediction=None):
        self.feature = feature
        self.threshold = threshold
        self.left = left
        self.right = right
        self.prediction = prediction

def train_tree(X, y, depth=0, max_depth=3):
    # Base case: pure leaf or max depth reached
    if calculate_gini(y) == 0 or depth == max_depth:
        return TreeNode(prediction=np.bincount(y).argmax())

    feat, thresh, gain = find_best_split(X, y)
    if feat is None or gain <= 0:
        return TreeNode(prediction=np.bincount(y).argmax())

    left_mask = X[:, feat] <= thresh
    left_node = train_tree(X[left_mask], y[left_mask], depth + 1, max_depth)
    right_node = train_tree(X[~left_mask], y[~left_mask], depth + 1, max_depth)

    return TreeNode(feature=feat, threshold=thresh, left=left_node, right=right_node)

# Sample dataset (sizes vs red (0) / blue (1))
X_demo = np.array([[2], [3], [4], [6], [7], [8], [9], [10]])
y_demo = np.array([0, 0, 0, 1, 1, 1, 1, 1])

root = train_tree(X_demo, y_demo, max_depth=2)
print(f"Optimal Root Split -> Feature Index: {root.feature}, Threshold: {root.threshold}")

2. Production Scikit-Learn with Hyperparameter Pruning

from sklearn.tree import DecisionTreeClassifier

# Guarding against overfitting with max_depth and min_samples_split
clf = DecisionTreeClassifier(
    criterion="gini",
    max_depth=3,
    min_samples_split=4,
    min_samples_leaf=2,
    random_state=42
)
clf.fit(X_demo, y_demo)

print("Scikit-Learn Root Feature:", clf.tree_.feature[0])
print("Scikit-Learn Root Threshold:", clf.tree_.threshold[0])

6. Overfitting Guardrails: 4 Hyperparameters to Control

• max_depth: Restricts the maximum vertical depth of the tree to prevent memorizing single outliers.
• min_samples_split: Sets the minimum number of samples required to consider splitting an internal node.
• min_samples_leaf: Guarantees that every leaf node contains at least $k$ samples.
• ccp_alpha (Cost-Complexity Pruning): Trims weak branches after full tree growth to maximize cross-validated generalization.

Summary

• Decision trees split datasets recursively by evaluating all possible feature thresholds.
• Gini Impurity measures classification disorder ($0.0$ is pure, $0.5$ is maximally mixed).
• Information Gain selects the split that maximizes the reduction in weighted child impurity.
• Regression Trees minimize Variance / MSE and output the sample mean $\bar{y}$ in leaves.
• Pre-pruning hyperparameters (max_depth, min_samples_leaf) are essential to prevent overfitting.