Skip to main content

Decision Trees with scikit-learn

This example trains a multiclass DecisionTreeClassifier on scikit-learn's bundled Iris dataset. The outer test set remains isolated while cross-validation selects both growth controls and the cost-complexity pruning parameter ccp_alpha.

1. Establish the evaluation boundary​

Iris contains 150 observations total—50 from each of three species. A stratified split preserves those class proportions. The tree does not need standardization: a monotonic rescaling changes threshold values but not the ordering of candidate splits.

2. Tune growth and pruning inside training folds​

The parameter grid deliberately compares pre-pruning (max_depth, min_samples_leaf) and post-pruning (ccp_alpha). GridSearchCV refits the selected configuration on all training data; only then is the untouched test set evaluated.

from sklearn.datasets import load_iris
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
from sklearn.model_selection import GridSearchCV, StratifiedKFold, train_test_split
from sklearn.tree import DecisionTreeClassifier, export_text

iris = load_iris()
X_train, X_test, y_train, y_test = train_test_split(
    iris.data,
    iris.target,
    test_size=0.25,
    random_state=42,
    stratify=iris.target,
)

search = GridSearchCV(
    DecisionTreeClassifier(random_state=42),
    param_grid={
        "criterion": ["gini", "entropy"],
        "max_depth": [2, 3, None],
        "min_samples_leaf": [1, 3, 5],
        "ccp_alpha": [0.0, 0.005, 0.01],
    },
    scoring="balanced_accuracy",
    cv=StratifiedKFold(n_splits=5, shuffle=True, random_state=42),
    n_jobs=1,
)
search.fit(X_train, y_train)

tree = search.best_estimator_
y_pred = tree.predict(X_test)
print("Best parameters:", search.best_params_)
print("Tree depth / leaves:", tree.get_depth(), tree.get_n_leaves())
print("Test accuracy:", round(accuracy_score(y_test, y_pred), 3))
print("Confusion matrix:\n", confusion_matrix(y_test, y_pred))
print(classification_report(y_test, y_pred, target_names=iris.target_names))
print(export_text(tree, feature_names=list(iris.feature_names)))

assert X_train.shape == (112, 4)
assert X_test.shape == (38, 4)
assert tree.get_n_leaves() >= 2
assert y_pred.shape == y_test.shape

3. Interpret the result carefully​

export_text exposes the exact rules, but that does not make every rule stable or causal. Nearby training samples can produce a different greedy tree, correlated features can substitute for one another, and the selected thresholds describe this fitted sample. Report depth and leaf count alongside predictive metrics, and check subgroup errors when the model affects people.

The classifier's feature_importances_ values are normalized totals of impurity decrease. They can favor continuous or high-cardinality features and split credit among correlated variables unpredictably. Treat them as properties of the fitted tree, not as causal effects; use held-out permutation importance when the question is predictive reliance.

4. Production considerations​

  • scikit-learn trees currently require numeric feature columns; encode categorical values deliberately rather than assigning arbitrary ordinal codes.
  • Missing-value support depends on estimator, splitter, and criterion. Verify the exact API contract instead of assuming every tree library uses surrogate splits.
  • If observations share a patient, customer, device, or time window, replace the stratified IID folds with group-aware or time-aware splits.
  • For class imbalance, choose a metric aligned with the cost of errors and consider class_weight; accuracy alone can hide minority-class failure.

This workflow keeps tuning inside training data, evaluates multiclass behavior directly, and distinguishes a pruned single tree from tree ensembles.