Implement a Decision Tree Split

Math Foundations

At each node $N$, a Decision Tree selects feature $j$ and continuous threshold $t$ to partition samples into left subset $L = {x \mid x_j \le t}$ and right subset $R = {x \mid x_j > t}$.

1. Gini Impurity

Measures the probability of incorrectly classifying a randomly chosen sample if labeled according to class distribution $p_k$:

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

• Pure Node (All samples belong to 1 class) $\implies \text{Gini} = 0.0$.
• Maximum Impurity ($50/50$ split in binary) $\implies \text{Gini} = 0.5$.

2. Information Gain (Gini Gain)

Measures impurity reduction from parent node $S$ to weighted child nodes $L$ and $R$:

$$\text{Gain}(S, j, t) = \text{Gini}(S) - \left( \frac{|L|}{|S|} \text{Gini}(L) + \frac{|R|}{|S|} \text{Gini}(R) \right)$$

The algorithm searches all features $j$ and thresholds $t$ to find:

$$(j^, t^) = \arg\max_{j, t} \text{Gain}(S, j, t)$$

Parent Node S (Gini = 0.48) ──► [ Split Feature X_j <= t ] ──┬──► Left Child L (Gini = 0.10)
                                                              └──► Right Child R (Gini = 0.05)

Python Implementation from Scratch

import numpy as np


class Node:

  def __init__(
      self, feature=None, threshold=None, left=None, right=None, value=None
  ):
    self.feature = feature
    self.threshold = threshold
    self.left = left
    self.right = right
    self.value = value  # Majority class for leaf nodes

  def is_leaf(self):
    return self.value is not None


class DecisionTreeFromScratch:

  def __init__(self, max_depth=5, min_samples_split=2):
    self.max_depth = max_depth
    self.min_samples_split = min_samples_split
    self.root = None

  def _gini(self, y):
    if len(y) == 0:
      return 0.0
    p = np.bincount(y) / len(y)
    return 1.0 - np.sum(p**2)

  def _best_split(self, X, y):
    best_gain = -1.0
    split_idx, split_thresh = None, None
    n_samples, n_features = X.shape
    parent_gini = self._gini(y)

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

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

        left_gini = self._gini(y[left_mask])
        right_gini = self._gini(y[right_mask])
        n_left, n_right = np.sum(left_mask), np.sum(right_mask)

        weighted_gini = (n_left / n_samples) * left_gini + (
            n_right / n_samples
        ) * right_gini
        gain = parent_gini - weighted_gini

        if gain > best_gain:
          best_gain = gain
          split_idx = feat_idx
          split_thresh = thresh

    return split_idx, split_thresh

  def _build_tree(self, X, y, depth=0):
    n_samples, n_classes = X.shape[0], len(np.unique(y))

    # Stopping criteria
    if (
        depth >= self.max_depth
        or n_classes == 1
        or n_samples < self.min_samples_split
    ):
      majority_class = np.argmax(np.bincount(y))
      return Node(value=majority_class)

    feat_idx, thresh = self._best_split(X, y)
    if feat_idx is None:
      majority_class = np.argmax(np.bincount(y))
      return Node(value=majority_class)

    left_mask = X[:, feat_idx] <= thresh
    left_child = self._build_tree(X[left_mask], y[left_mask], depth + 1)
    right_child = self._build_tree(X[~left_mask], y[~left_mask], depth + 1)

    return Node(
        feature=feat_idx,
        threshold=thresh,
        left=left_child,
        right=right_child,
    )

  def fit(self, X, y):
    self.root = self._build_tree(X, y)

  def _traverse(self, x, node):
    if node.is_leaf():
      return node.value
    if x[node.feature] <= node.threshold:
      return self._traverse(x, node.left)
    return self._traverse(x, node.right)

  def predict(self, X):
    return np.array([self._traverse(x, self.root) for x in X])

Say this out loud

Implementing a decision tree split evaluates Gini Impurity across all candidate feature thresholds. The greedy algorithm selects the feature and threshold that maximizes Information Gain. Tree recursion continues until max depth, pure nodes, or minimum sample thresholds stop growth to prevent overfitting.

Followups to expect

1. What is Entropy vs Gini Impurity? Entropy uses logarithmic measure $- \sum p_i \log_2(p_i)$, which is computationally slower than Gini Impurity $1 - \sum p_i^2$, though both produce nearly identical tree splits in practice.
2. How do decision trees handle missing values during splits? Algorithms like XGBoost assign a default split direction (left or right) for missing values at each node, learning the optimal missing direction during training.