Decision Tree (CART)

"""CART classification tree - the split criterion, transcribed.

For every feature and every candidate threshold, split, measure the weighted
impurity of the two children, and keep the best. That exhaustive scan is the
objective above, written out literally: nothing is cached, nothing is sorted,
and every candidate recomputes both children from scratch.
"""


def gini(labels):
    """H(S) = 1 - sum_c p_c^2"""
    if not labels:
        return 0.0

    counts = {}
    for label in labels:
        counts[label] = counts.get(label, 0) + 1

    impurity = 1.0
    for count in counts.values():
        proportion = count / len(labels)
        impurity -= proportion * proportion
    return impurity


def majority(labels):
    counts = {}
    for label in labels:
        counts[label] = counts.get(label, 0) + 1

    best_label = None
    best_count = -1
    for label, count in counts.items():
        if count > best_count:
            best_label = label
            best_count = count
    return best_label


def best_split(X, y):
    """Exhaustive search over every feature and every midpoint."""
    n = len(X)
    d = len(X[0])
    best = (None, None, gini(y))

    for feature in range(d):
        values = sorted({X[i][feature] for i in range(n)})

        # Candidate thresholds are midpoints between adjacent distinct values.
        for position in range(len(values) - 1):
            threshold = (values[position] + values[position + 1]) / 2.0

            left_labels = []
            right_labels = []
            for i in range(n):
                if X[i][feature] <= threshold:
                    left_labels.append(y[i])
                else:
                    right_labels.append(y[i])

            if not left_labels or not right_labels:
                continue

            # Weighted by child size, which is what stops the search from
            # preferring a split that isolates a single point.
            weighted = (
                len(left_labels) / n * gini(left_labels)
                + len(right_labels) / n * gini(right_labels)
            )

            if weighted < best[2]:
                best = (feature, threshold, weighted)

    return best


def build(X, y, depth=0, max_depth=5, min_samples_leaf=1):
    if depth >= max_depth or len(set(y)) == 1 or len(y) < 2 * min_samples_leaf:
        return {"leaf": True, "prediction": majority(y)}

    feature, threshold, _ = best_split(X, y)
    if feature is None:
        return {"leaf": True, "prediction": majority(y)}

    left_X, left_y, right_X, right_y = [], [], [], []
    for i in range(len(X)):
        if X[i][feature] <= threshold:
            left_X.append(X[i])
            left_y.append(y[i])
        else:
            right_X.append(X[i])
            right_y.append(y[i])

    return {
        "leaf": False,
        "feature": feature,
        "threshold": threshold,
        "left": build(left_X, left_y, depth + 1, max_depth, min_samples_leaf),
        "right": build(right_X, right_y, depth + 1, max_depth, min_samples_leaf),
    }


def predict_one(node, x):
    while not node["leaf"]:
        node = node["left"] if x[node["feature"]] <= node["threshold"] else node["right"]
    return node["prediction"]

"""CART - typed, presorted, incremental class counts."""

from dataclasses import dataclass

import numpy as np
from numpy.typing import NDArray

Vector = NDArray[np.float64]
Matrix = NDArray[np.float64]


@dataclass(frozen=True)
class Split:
    feature: int
    threshold: float
    impurity: float


@dataclass(frozen=True)
class Node:
    """A leaf carries a prediction; an internal node carries a split.

    Modelled as one frozen type with optional children rather than a dict,
    so a malformed node is a construction error rather than a KeyError during
    prediction.
    """

    prediction: int
    split: Split | None = None
    left: "Node | None" = None
    right: "Node | None" = None

    @property
    def is_leaf(self) -> bool:
        return self.split is None


def _best_split(X: Matrix, y: NDArray, n_classes: int, min_samples_leaf: int) -> Split | None:
    """Scan each feature in sorted order, moving one point at a time.

    The key change: sorting the feature once means the class counts for the two
    children can be updated incrementally as the boundary sweeps right, instead
    of being recomputed for every candidate threshold.
    """
    n, d = X.shape
    parent_counts = np.bincount(y, minlength=n_classes)
    best: Split | None = None
    best_impurity = 1.0 - float((parent_counts / n) ** 2 @ np.ones(n_classes))

    for feature in range(d):
        order = np.argsort(X[:, feature], kind="stable")
        values = X[order, feature]
        labels = y[order]

        left_counts = np.zeros(n_classes, dtype=np.int64)
        right_counts = parent_counts.copy()

        for position in range(n - 1):
            label = labels[position]
            left_counts[label] += 1
            right_counts[label] -= 1

            # Skip boundaries inside a run of equal values: they are not
            # separable, and splitting there produces two identical children.
            if values[position] == values[position + 1]:
                continue

            n_left = position + 1
            n_right = n - n_left
            if n_left < min_samples_leaf or n_right < min_samples_leaf:
                continue

            gini_left = 1.0 - float(((left_counts / n_left) ** 2).sum())
            gini_right = 1.0 - float(((right_counts / n_right) ** 2).sum())
            weighted = (n_left * gini_left + n_right * gini_right) / n

            if weighted < best_impurity:
                best_impurity = weighted
                best = Split(
                    feature=feature,
                    threshold=float((values[position] + values[position + 1]) / 2.0),
                    impurity=weighted,
                )

    return best


def build(
    X: Matrix,
    y: NDArray,
    n_classes: int,
    depth: int = 0,
    max_depth: int = 5,
    min_samples_leaf: int = 1,
) -> Node:
    """Grow a tree. Raises ValueError on malformed input."""
    if X.ndim != 2:
        raise ValueError(f"X must be 2-D, got shape {X.shape}")
    if X.shape[0] != y.shape[0]:
        raise ValueError(f"X has {X.shape[0]} rows but y has {y.shape[0]}")
    if min_samples_leaf < 1:
        raise ValueError(f"min_samples_leaf must be at least 1, got {min_samples_leaf}")

    prediction = int(np.bincount(y, minlength=n_classes).argmax())

    if depth >= max_depth or np.unique(y).size == 1 or y.size < 2 * min_samples_leaf:
        return Node(prediction=prediction)

    split = _best_split(X, y, n_classes, min_samples_leaf)
    if split is None:
        return Node(prediction=prediction)

    mask = X[:, split.feature] <= split.threshold
    return Node(
        prediction=prediction,
        split=split,
        left=build(X[mask], y[mask], n_classes, depth + 1, max_depth, min_samples_leaf),
        right=build(X[~mask], y[~mask], n_classes, depth + 1, max_depth, min_samples_leaf),
    )


def predict(node: Node, X: Matrix) -> NDArray:
    out = np.empty(X.shape[0], dtype=np.int64)
    for index, row in enumerate(X):
        current = node
        while not current.is_leaf:
            assert current.split is not None
            current = (
                current.left if row[current.split.feature] <= current.split.threshold
                else current.right
            )
            if current is None:
                raise ValueError("malformed tree: internal node with a missing child")
        out[index] = current.prediction
    return out

"""CART - histogram binning, one pass per feature, preallocated buffers."""

import numpy as np
from numpy.typing import NDArray

Matrix = NDArray[np.float64]


class HistogramTreeBuilder:
    """Split finding by binned histograms - what every production tree does.

    Features are quantized ONCE into at most 256 bins, before any tree is
    grown. Finding a split then means accumulating a per-bin class histogram in
    one linear pass and scanning 255 candidate boundaries, instead of sorting
    and sweeping n values. The per-node cost stops depending on n except through
    the single accumulation pass.

    The trade is exact for what it is: splits are now chosen from bin edges
    rather than from every midpoint, so the fitted thresholds are approximate.
    On real data the accuracy cost is tiny and the speedup is large, which is
    why this is the default - but on a feature whose signal lives in a narrow
    range, the binning can quantize the split away entirely.
    """

    N_BINS = 256

    def __init__(self, X: Matrix, y: NDArray, n_classes: int) -> None:
        self._n_classes = n_classes
        self._y = np.ascontiguousarray(y, dtype=np.int64)

        # Quantile edges, computed once for the whole tree - and reusable
        # across every tree in an ensemble, which is where this really pays.
        design = np.ascontiguousarray(X, dtype=np.float64)
        self._edges = np.stack([
            np.quantile(design[:, feature], np.linspace(0, 1, self.N_BINS + 1)[1:-1])
            for feature in range(design.shape[1])
        ])

        # uint8 binned matrix: 8x smaller than float64, so far more of it stays
        # in cache during the accumulation pass.
        self._binned = np.empty(design.shape, dtype=np.uint8, order="F")
        for feature in range(design.shape[1]):
            np.searchsorted(
                self._edges[feature], design[:, feature], side="left",
            ).astype(np.uint8, copy=False, casting="unsafe")
            self._binned[:, feature] = np.searchsorted(
                self._edges[feature], design[:, feature], side="left"
            )

        # Histogram buffer, allocated once and reused at every node.
        self._histogram = np.empty((self.N_BINS, n_classes), dtype=np.int64)

    def best_split(self, indices: NDArray, min_samples_leaf: int) -> tuple[int, int, float]:
        """Returns (feature, bin, impurity) for the best split of this node."""
        n = indices.size
        labels = self._y[indices]
        best = (-1, -1, np.inf)

        for feature in range(self._binned.shape[1]):
            # One pass over the node's rows builds the whole histogram: bincount
            # over the flattened (bin, class) index, no Python-level loop.
            bins = self._binned[indices, feature].astype(np.int64, copy=False)
            flat = bins * self._n_classes + labels
            counts = np.bincount(flat, minlength=self.N_BINS * self._n_classes)
            self._histogram[:] = counts.reshape(self.N_BINS, self._n_classes)

            # Prefix sums turn "counts per bin" into "counts left of each
            # boundary" for all 255 candidates at once.
            left = np.cumsum(self._histogram, axis=0)
            total = left[-1]
            right = total - left

            n_left = left.sum(axis=1)
            n_right = n - n_left
            valid = (n_left >= min_samples_leaf) & (n_right >= min_samples_leaf)
            if not valid.any():
                continue

            with np.errstate(invalid="ignore", divide="ignore"):
                gini_left = 1.0 - ((left / n_left[:, None]) ** 2).sum(axis=1)
                gini_right = 1.0 - ((right / n_right[:, None]) ** 2).sum(axis=1)
                weighted = (n_left * gini_left + n_right * gini_right) / n

            weighted[~valid] = np.inf
            candidate = int(weighted.argmin())
            if weighted[candidate] < best[2]:
                best = (feature, candidate, float(weighted[candidate]))

        return best