Decision Trees

Plot each dataset before choosing a tree strategy; the task is a collection of different geometries, not one shared distribution.

Check feature scale, covariance, and orientation. Can a data-driven rotation make the tree’s axis-aligned splits more useful?

Use the unlabeled test points to stabilize transductive feature preprocessing, while keeping their labels out of training and model selection.

Select pruning strength with stratified CV, and be cautious about tiny CV-score differences on small datasets.

Symmetrically whiten pooled features, rotate them to reduce 24-bin histogram-estimated mutual information, then select pruning by repeated stratified CV with a prior based on the CV-best alpha; measured result: 66.0/100 on released collection A.

"""Decision-tree solution for the Polish 2026 Stage 2 task.

The input consists of many independent two-feature binary classification
problems.  We use the unlabeled test points to estimate a stable whitening
transform, rotate the axes to reduce dependence between the training
projections, then choose the tree pruning strength by repeated stratified CV.
"""

from __future__ import annotations

import numpy as np
from sklearn.model_selection import RepeatedStratifiedKFold
from sklearn.tree import DecisionTreeClassifier


def preprocess_data(X_train: np.ndarray, X_test: np.ndarray):
    """Whiten pooled features and choose a low-mutual-information rotation.

    The histogram objective is estimated from the pooled, unlabeled point
    cloud. Both train and test points share the resulting transform. The function keeps
    the required two-column shape and has a safe identity fallback for invalid
    or unexpected inputs.
    """
    X_train = np.asarray(X_train, dtype=np.float64)
    X_test = np.asarray(X_test, dtype=np.float64)

    if (
        X_train.ndim != 2
        or X_test.ndim != 2
        or X_train.shape[1] != 2
        or X_test.shape[1] != 2
        or len(X_train) < 2
        or len(X_test) == 0
    ):
        return X_train.copy(), X_test.copy()

    pooled = np.concatenate((X_train, X_test), axis=0)
    if not np.isfinite(pooled).all():
        return X_train.copy(), X_test.copy()

    center = pooled.mean(axis=0)
    covariance = np.cov(pooled, rowvar=False)
    eigenvalues, eigenvectors = np.linalg.eigh(covariance)
    floor = max(float(eigenvalues.max()), 1.0) * 1e-12

    # Symmetric whitening preserves the original coordinate frame while
    # correcting for scale and correlation.
    whitening = (
        eigenvectors
        @ np.diag(1.0 / np.sqrt(np.maximum(eigenvalues, floor)))
        @ eigenvectors.T
    )
    Z = (pooled - center) @ whitening

    n_train = len(X_train)
    angles = np.deg2rad(np.arange(180, dtype=np.float64))
    cosines = np.cos(angles)
    sines = np.sin(angles)
    first = Z[:, 0, None] * cosines + Z[:, 1, None] * sines
    second = -Z[:, 0, None] * sines + Z[:, 1, None] * cosines

    # Estimate mutual information from 24 equal-frequency bins on each axis.
    # The small pseudocount makes the estimate finite for sparse histograms.
    mutual_information = np.empty(len(angles), dtype=np.float64)
    quantiles = np.linspace(0.0, 1.0, 25)
    for j in range(len(angles)):
        edges_1 = np.unique(np.quantile(first[:, j], quantiles))
        edges_2 = np.unique(np.quantile(second[:, j], quantiles))
        if len(edges_1) < 2 or len(edges_2) < 2:
            mutual_information[j] = np.inf
            continue

        counts = np.histogram2d(
            first[:, j], second[:, j], bins=(edges_1, edges_2)
        )[0]
        counts += 1e-6
        joint = counts / counts.sum()
        marginal_1 = joint.sum(axis=1, keepdims=True)
        marginal_2 = joint.sum(axis=0, keepdims=True)
        mutual_information[j] = np.sum(
            joint * np.log(joint / (marginal_1 @ marginal_2))
        )

    best = int(np.argmin(mutual_information))
    c, s = cosines[best], sines[best]
    rotation = np.array([[c, -s], [s, c]])
    Z = Z @ rotation
    return Z[:n_train], Z[n_train:]


def get_decision_tree_hyperparameters(X_train: np.ndarray, y_train: np.ndarray):
    """Select cost-complexity pruning strength by repeated stratified CV."""
    X_train = np.asarray(X_train, dtype=np.float64)
    y_train = np.asarray(y_train).ravel()
    classes, counts = np.unique(y_train, return_counts=True)
    fallback_alpha = 0.0015

    if len(classes) < 2:
        return {"ccp_alpha": fallback_alpha, "random_state": 42}

    n_splits = min(4, int(counts.min()))
    if n_splits < 2:
        return {"ccp_alpha": fallback_alpha, "random_state": 42}

    alphas = np.round(0.0015 + 0.00025 * np.arange(13), 5)
    cv = RepeatedStratifiedKFold(
        n_splits=n_splits, n_repeats=2, random_state=1741
    )
    folds = list(cv.split(X_train, y_train))
    scores = np.zeros((len(alphas), len(folds)), dtype=np.float64)

    for alpha_index, alpha in enumerate(alphas):
        for fold_index, (train_index, valid_index) in enumerate(folds):
            tree = DecisionTreeClassifier(
                ccp_alpha=float(alpha), random_state=42
            )
            tree.fit(X_train[train_index], y_train[train_index])
            scores[alpha_index, fold_index] = tree.score(
                X_train[valid_index], y_train[valid_index]
            )

    mean_scores = scores.mean(axis=1)
    best_score = float(mean_scores.max())

    # The selected global policy uses a prior conditioned on the best CV alpha.
    # Both values are fixed across tasks; the condition depends only on this
    # dataset's training labels and repeated-CV scores.
    best_index = int(np.argmax(mean_scores))
    cv_best_alpha = float(alphas[best_index])
    prior_alpha = 0.0015 if cv_best_alpha <= 0.003 else 0.00375

    # Retain near-best candidates because individual CV estimates are noisy.
    eligible = np.flatnonzero(mean_scores >= best_score - 0.001)
    selected = min(
        eligible, key=lambda index: abs(float(alphas[index]) - prior_alpha)
    )
    return {"ccp_alpha": float(alphas[selected]), "random_state": 42}


def fit_tree(X_train: np.ndarray, y_train: np.ndarray) -> DecisionTreeClassifier:
    """Convenience helper for local use of the two task functions above."""
    params = get_decision_tree_hyperparameters(X_train, y_train)
    return DecisionTreeClassifier(**params).fit(X_train, y_train)

Download solution.py