S089 · Decision Trees — Gini, Entropy, Splits

# tree_scratch.py — decision tree from scratch + sklearn comparison, ~140 lines.
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
 
np.random.seed(42)
 
# --- Data ---
X, y = load_iris(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, stratify=y, random_state=42)
 
# =====================================================
# 1. Impurity functions
# =====================================================
def gini(y):
    if len(y) == 0:
        return 0.0
    _, counts = np.unique(y, return_counts=True)
    p = counts / len(y)
    return 1.0 - (p ** 2).sum()
 
def entropy(y):
    if len(y) == 0:
        return 0.0
    _, counts = np.unique(y, return_counts=True)
    p = counts / len(y)
    return -(p * np.log2(p + 1e-12)).sum()
 
# =====================================================
# 2. Best split — brute force
# =====================================================
def best_split(X, y, impurity=gini):
    """Return (feature_idx, threshold, impurity_gain) for the best split, or None."""
    n, d = X.shape
    parent_imp = impurity(y)
    best = None
    for j in range(d):
        vals = np.unique(X[:, j])
        for t in (vals[:-1] + vals[1:]) / 2:            # midpoints between consecutive values
            left = X[:, j] < t
            right = ~left
            n_l, n_r = left.sum(), right.sum()
            if n_l == 0 or n_r == 0:
                continue
            weighted_imp = (n_l * impurity(y[left]) + n_r * impurity(y[right])) / n
            gain = parent_imp - weighted_imp
            if best is None or gain > best[2]:
                best = (j, t, gain)
    return best
 
# =====================================================
# 3. Tree node
# =====================================================
class Node:
    def __init__(self):
        self.feature = None      # split feature index
        self.threshold = None    # split threshold
        self.left = None
        self.right = None
        self.prediction = None   # for leaves: majority class
 
def build(X, y, depth=0, max_depth=5, min_samples=2, impurity=gini):
    node = Node()
    # Stopping criteria
    if depth >= max_depth or len(y) < min_samples or len(np.unique(y)) == 1:
        vals, counts = np.unique(y, return_counts=True)
        node.prediction = vals[counts.argmax()]
        return node
    split = best_split(X, y, impurity)
    if split is None or split[2] <= 0:
        vals, counts = np.unique(y, return_counts=True)
        node.prediction = vals[counts.argmax()]
        return node
    j, t, _ = split
    node.feature, node.threshold = j, t
    mask = X[:, j] < t
    node.left = build(X[mask], y[mask], depth + 1, max_depth, min_samples, impurity)
    node.right = build(X[~mask], y[~mask], depth + 1, max_depth, min_samples, impurity)
    return node
 
def predict_one(node, x):
    if node.prediction is not None:
        return node.prediction
    if x[node.feature] < node.threshold:
        return predict_one(node.left, x)
    return predict_one(node.right, x)
 
def predict(root, X):
    return np.array([predict_one(root, x) for x in X])
 
# =====================================================
# 4. Fit + evaluate
# =====================================================
for max_depth in [1, 2, 3, 5, 10, None]:
    md = max_depth if max_depth else 20
    root = build(X_train, y_train, max_depth=md)
    acc_tr = accuracy_score(y_train, predict(root, X_train))
    acc_te = accuracy_score(y_test, predict(root, X_test))
    print(f"[scratch · depth={str(max_depth):>4}]  train={acc_tr:.3f}  test={acc_te:.3f}")
 
# =====================================================
# 5. sklearn sanity check
# =====================================================
print()
for max_depth in [1, 2, 3, 5, 10, None]:
    sk = DecisionTreeClassifier(max_depth=max_depth, random_state=42).fit(X_train, y_train)
    acc_tr = accuracy_score(y_train, sk.predict(X_train))
    acc_te = accuracy_score(y_test, sk.predict(X_test))
    print(f"[sklearn · depth={str(max_depth):>4}]  train={acc_tr:.3f}  test={acc_te:.3f}")
 
# =====================================================
# 6. Print the fitted tree (small enough to read)
# =====================================================
def print_tree(node, feature_names, class_names, indent=""):
    if node.prediction is not None:
        print(f"{indent}→ predict {class_names[node.prediction]}")
        return
    print(f"{indent}if {feature_names[node.feature]} < {node.threshold:.2f}:")
    print_tree(node.left, feature_names, class_names, indent + "  ")
    print(f"{indent}else:")
    print_tree(node.right, feature_names, class_names, indent + "  ")
 
print("\n=== Learned tree (max_depth=3) ===")
root3 = build(X_train, y_train, max_depth=3)
print_tree(root3, ["sepal_len", "sepal_wid", "petal_len", "petal_wid"], ["setosa", "versicolor", "virginica"])