Decision Trees — CART from Scratch

Python · scikit-learn · R (rpart)  ·  Download CART module

Splits, Purity, and Pruning

A decision tree splits the feature space with axis-aligned cuts, each chosen to make the two child nodes as pure as possible — Gini or entropy for classification, variance for regression. Searching every feature and every threshold sounds expensive; sorting each feature once and sweeping the cut point with cumulative counts makes it O(pnlogn)O(p\,n\log n) per node. The data are 30,000 credit-card clients with a 22.1% default rate, and the tree's first question is the one intuition expects: PAY_1 ≤ 1.5 — was the client already behind last month.

Δi(s,t)=i(t)pLi(tL)pRi(tR),Rα(T)=R(T)+αleaves(T)\Delta i(s,t) = i(t) - p_L\,i(t_L) - p_R\,i(t_R), \qquad R_\alpha(T) = R(T) + \alpha\,|\text{leaves}(T)|

Does it match scikit-learn?

The from-scratch implementation is checked against scikit-learn, and the agreement is worth stating precisely rather than loosely. At depths 2 and 3 the trees are bit-identical — same splits, same leaves, 100% identical test labels. Agreement then degrades gracefully: 99.99% at depth 5, 99.31% at depth 8. The cause is tie-breaking, not a difference of algorithm: when two splits give exactly equal impurity decrease, scikit-learn scans features in randomised order and this implementation scans them in index order, so the tie falls the other way and the subtree below it diverges. Accuracy and AUC still match to 3–4 decimals.

depthAUC from scratchAUC sklearnleavesidentical labels
20.68980.68984 / 4100%
30.72590.72598 / 8100%
50.74630.746231 / 3199.99%
80.73620.7365152 / 15199.31%

Cost-complexity pruning

Grown to purity a tree overfits — here to 1,804 leaves and a test AUC of 0.6535. Breiman's remedy is cost-complexity pruning: penalise the tree by its number of leaves and collapse the weakest link first, giving a nested sequence of subtrees indexed by α\alpha, from which cross-validation picks one — splitting the training rows into kk folds, fitting on k1k-1 and scoring on the one held out, then averaging, so every candidate α\alpha is judged on data it did not see and the test set stays untouched until the end (Model Selection sets it beside AIC/BIC and Bayesian LOO). It prunes to 7 leaves and lifts test AUC to 0.7443 — a smaller tree that generalises better, which is the bias–variance trade-off made concrete.

The pruning paths differ slightly in length (895 alphas against scikit-learn's 920) while spanning the identical α\alpha range. That is inherited rather than a different rule: the two fully-grown trees already differ by a handful of leaves through the tie-breaking above. Pruned at a shared α\alpha the gap closes as α\alpha grows — 1148 against 1131 leaves while the tree is still large, and exactly equal (11 vs 11) by the time α\alpha has pruned to the size cross-validation selects.

Why one tree is never enough

The reason this example opens an arc rather than closing one: a single deep tree pins a given client's default probability only to an SD of about 0.26 across resamples, with test AUC 0.651 ± 0.006. Averaging 25 such trees lifts AUC to 0.758 — the variance cancels while the signal adds. That is bagging, and every later example in this section is a refinement of it.

Against a logistic regression

Because this collection is otherwise econometric, the tree is benchmarked against a logistic regression on the same split. The tree wins, but not by much: 0.744 against 0.715, with a random forest at 0.762. The logit buys something the tree cannot state — each extra month in arrears multiplies the odds of default by 2.0, the one-line summary a scorecard needs — while the tree buys flexibility, stepping to the kink where default is flat for paid-up clients and jumps once in arrears. The ordering is also less robust than it looks: in the R notebook, where rpart prunes to its 1-SE default rather than the cross-validated optimum, the logit wins instead (0.712 against 0.688). "The tree beats the regression" is a statement about a tuned tree.

Notebooks

Downloads

CART Module — Source Code

"""
cart.py -- Classification And Regression Trees (CART) from scratch (Breiman, Friedman, Olshen & Stone 1984).

Backs the notebooks in  Z-ML Trees-Decision Trees (CART).

A decision tree recursively splits the feature space with axis-aligned cuts, each chosen to make the two
child nodes as PURE as possible. Purity is measured by an impurity criterion:

    Gini     i(t) = 1 - sum_k p_k^2            (classification; default)
    Entropy  i(t) = -sum_k p_k log p_k          (classification)
    MSE      i(t) = (1/n) sum (y - ybar)^2      (regression = variance)

A split s sends a fraction p_L of the node's samples left and p_R right; its quality is the impurity
DECREASE  di = i(t) - p_L i(t_L) - p_R i(t_R), and at each node we search every feature and every
threshold for the split that maximises it. The best split per feature is found in O(n log n) by sorting
the feature once and sweeping the cut point with cumulative class counts (classification) or cumulative
sums (regression), so the whole search is O(n_features * n log n).

Grown to purity a tree OVERFITS. CART's remedy is COST-COMPLEXITY PRUNING: penalise the tree by the number
of leaves, R_alpha(T) = R(T) + alpha|leaves(T)|, and for increasing alpha collapse the "weakest-link"
subtree first -- the internal node t whose effective penalty
        alpha_eff(t) = (R(t) - R(T_t)) / (|leaves(T_t)| - 1)
is smallest. This yields a nested sequence of subtrees indexed by alpha; the right alpha is picked by
cross-validation. R(t) here is the node impurity weighted by the share of samples reaching it, matching
scikit-learn's `ccp_alpha`.
"""

import numpy as np


class _Node:
    __slots__ = ("feat", "thr", "left", "right", "value", "n", "imp")
    def __init__(self, value, n, imp):
        self.feat = None; self.thr = None; self.left = None; self.right = None
        self.value = value            # leaf prediction: class-prob vector (clf) or mean (reg)
        self.n = n; self.imp = imp    # samples reaching node; node impurity


class DecisionTree:
    def __init__(self, criterion="gini", max_depth=None, min_samples_split=2,
                 min_samples_leaf=1, ccp_alpha=0.0):
        self.criterion = criterion; self.max_depth = max_depth
        self.min_samples_split = min_samples_split; self.min_samples_leaf = min_samples_leaf
        self.ccp_alpha = ccp_alpha
        self.regression = (criterion == "mse")

    # ---- impurity of a set of labels ----
    def _impurity(self, y):
        if self.regression:
            return np.mean((y - y.mean()) ** 2) if len(y) else 0.0
        cnt = np.bincount(y, minlength=self.n_classes_).astype(float); p = cnt / cnt.sum()
        if self.criterion == "entropy":
            p = p[p > 0]; return float(-(p * np.log2(p)).sum())
        return float(1.0 - (p * p).sum())                                     # gini

    def _leaf_value(self, y):
        if self.regression: return float(y.mean())
        return np.bincount(y, minlength=self.n_classes_).astype(float) / len(y)

    # ---- best split over all features (cumulative sweep) ----
    def _best_split(self, X, y):
        n, p = X.shape; parent_imp = self._impurity(y)
        best = (0.0, None, None)                                              # (gain, feat, thr)
        for j in range(p):
            xs = X[:, j]; order = np.argsort(xs, kind="mergesort")
            xj = xs[order]; yj = y[order]
            cut = np.where(np.diff(xj) > 0)[0]                                 # positions with a real threshold
            if len(cut) == 0: continue
            if self.regression:
                cs = np.cumsum(yj); cs2 = np.cumsum(yj * yj); tot = cs[-1]; tot2 = cs2[-1]
                nL = cut + 1.0; nR = n - nL
                sseL = cs2[cut] - cs[cut] ** 2 / nL
                sseR = (tot2 - cs2[cut]) - (tot - cs[cut]) ** 2 / nR
                child = (sseL + sseR) / n                                      # weighted child MSE
                gain = parent_imp - child
            else:
                oh = np.zeros((n, self.n_classes_)); oh[np.arange(n), yj] = 1.0
                cc = np.cumsum(oh, axis=0)                                     # cumulative class counts
                nL = (cut + 1.0)[:, None]; nR = n - nL
                cL = cc[cut]; cR = cc[-1] - cL
                giniL = 1.0 - ((cL / nL) ** 2).sum(1); giniR = 1.0 - ((cR / nR) ** 2).sum(1)
                if self.criterion == "entropy":
                    pL = cL / nL; pR = cR / nR
                    giniL = -np.where(pL > 0, pL * np.log2(np.where(pL > 0, pL, 1)), 0).sum(1)
                    giniR = -np.where(pR > 0, pR * np.log2(np.where(pR > 0, pR, 1)), 0).sum(1)
                child = (nL[:, 0] * giniL + nR[:, 0] * giniR) / n
                gain = parent_imp - child
            ok = (nL.reshape(-1) >= self.min_samples_leaf) & (nR.reshape(-1) >= self.min_samples_leaf)
            if not ok.any(): continue
            gain = np.where(ok, gain, -np.inf); k = int(np.argmax(gain))
            if gain[k] > best[0]:
                best = (float(gain[k]), j, float((xj[cut[k]] + xj[cut[k] + 1]) / 2.0))
        return best

    def _build(self, X, y, depth):
        node = _Node(self._leaf_value(y), len(y), self._impurity(y))
        if (len(y) < self.min_samples_split or node.imp == 0.0 or
                (self.max_depth is not None and depth >= self.max_depth)):
            return node
        gain, feat, thr = self._best_split(X, y)
        if feat is None or gain <= 0:
            return node
        m = X[:, feat] <= thr
        node.feat = feat; node.thr = thr
        node.left = self._build(X[m], y[m], depth + 1)
        node.right = self._build(X[~m], y[~m], depth + 1)
        return node

    def fit(self, X, y):
        X = np.asarray(X, float); y = np.asarray(y)
        self.n_features_ = X.shape[1]
        if not self.regression:
            self.classes_, y = np.unique(y, return_inverse=True); self.n_classes_ = len(self.classes_)
        self.n_total_ = len(y)
        self.root = self._build(X, y, 0)
        if self.ccp_alpha > 0: self._prune(self.ccp_alpha)
        return self

    # ---- prediction ----
    def _leaf(self, x, node):
        while node.feat is not None:
            node = node.left if x[node.feat] <= node.thr else node.right
        return node.value

    def predict_proba(self, X):
        X = np.asarray(X, float); return np.array([self._leaf(x, self.root) for x in X])

    def predict(self, X):
        if self.regression:
            X = np.asarray(X, float); return np.array([self._leaf(x, self.root) for x in X])
        return self.classes_[np.argmax(self.predict_proba(X), axis=1)]

    # ---- cost-complexity pruning ----
    def _R(self, node):                                                       # weighted resubstitution risk of a node
        return node.imp * node.n / self.n_total_

    def _collect(self, node, out):                                            # internal nodes with (RT, nleaves)
        if node.feat is None: return self._R(node), 1
        RTl, nl = self._collect(node.left, out); RTr, nr = self._collect(node.right, out)
        RT = RTl + RTr; nleaves = nl + nr
        out.append((node, self._R(node), RT, nleaves))
        return RT, nleaves

    def cost_complexity_pruning_path(self):
        """weakest-link sequence: increasing effective alphas at which each subtree collapses."""
        alphas = [0.0]
        import copy; root = copy.deepcopy(self.root); saved = self.root
        self.root = root
        while root.feat is not None:
            nodes = []; self._collect(root, nodes)
            nd, R, RT, nl = min(nodes, key=lambda z: (z[1] - z[2]) / (z[3] - 1))
            alphas.append((R - RT) / (nl - 1)); nd.feat = None; nd.left = nd.right = None
        self.root = saved
        return np.array(alphas)

    def _prune(self, alpha):
        while self.root.feat is not None:
            nodes = []; self._collect(self.root, nodes)
            nd, R, RT, nl = min(nodes, key=lambda z: (z[1] - z[2]) / (z[3] - 1))
            if (R - RT) / (nl - 1) > alpha: break
            nd.feat = None; nd.left = nd.right = None

    # ---- introspection ----
    def n_leaves(self):
        def rec(nd): return 1 if nd.feat is None else rec(nd.left) + rec(nd.right)
        return rec(self.root)

    def get_depth(self):
        def rec(nd, d): return d if nd.feat is None else max(rec(nd.left, d + 1), rec(nd.right, d + 1))
        return rec(self.root, 0)

    @property
    def feature_importances_(self):
        imp = np.zeros(getattr(self, "n_features_", self._infer_p()))
        def rec(nd):
            if nd.feat is None: return
            dec = nd.imp * nd.n - nd.left.imp * nd.left.n - nd.right.imp * nd.right.n
            imp[nd.feat] += dec; rec(nd.left); rec(nd.right)
        rec(self.root); s = imp.sum()
        return imp / s if s > 0 else imp

    def _infer_p(self):
        p = [0]
        def rec(nd):
            if nd.feat is None: return
            p[0] = max(p[0], nd.feat); rec(nd.left); rec(nd.right)
        rec(self.root); return p[0] + 1

References