Decision trees: twenty questions with data

Remember the game of twenty questions? "Is it alive?" "Is it bigger than a cat?" A good player does not ask random questions; they ask the one that cuts the remaining possibilities roughly in half. A decision tree is a model that plays this game with data — and because its questions are written in plain language, it is one of the few models you can read like a flowchart.

Try this

• Set the max depth slider to 1. The tree may ask only one question. Read it in the list of splits: which feature did the algorithm consider most informative?
• Raise the max depth slider one step at a time. Watch the training accuracy climb and compare it with the test accuracy at each depth.
• Push the max depth slider to the maximum. Is the training accuracy close to perfect? What happened to the test accuracy?
• Find the depth where test accuracy is best and read the top splits there. Could you explain this tree to someone in two sentences?

Now that you've tried it, let's see what was actually going on.

The idea in plain words

A decision tree predicts by asking a sequence of yes-or-no questions about the features: "Is the passenger's age below some value?", "Did they travel in first class?". Each answer sends the example down one branch, and the next question depends on the previous answers. At the end of every path sits a leaf with a prediction — usually the most common label among the training examples that ended up there.

In this lesson we use data on 891 passengers of the Titanic, of whom 342 survived. The features include ticket class, sex, age, the number of relatives on board and the fare paid. The tree learns which questions about these features best separate those who survived from those who did not.

Think of a triage nurse in an emergency room. She does not compute a weighted sum; she follows a short chain of questions — breathing? bleeding? conscious? — each one narrowing down what to do. A tree is that chain, learned from data instead of written by doctors.

Why it works

The tree is built from the top, one question at a time. At each step the algorithm tries every feature and every possible cut-off value and asks: which question splits this group into two groups that are as pure as possible, meaning each contains mostly one label? A popular measure of impurity is the Gini index of a group:

Gini = 1 − Σ pₖ²

where pₖ is the share of label k in the group. A group where everyone has the same label has Gini 0; a group split half and half between two labels has Gini 0.5, the worst possible for two classes. The algorithm picks the question that lowers the average impurity the most, then repeats the procedure separately inside each branch. This greedy, step-by-step search does not guarantee the best possible tree, but it is fast and works well in practice.

Left alone, the tree keeps splitting until every leaf is pure — even if a leaf holds a single passenger. At that point it has memorised the training data, noise and all: the overfitting of lesson 7 in a new costume. The main control is the maximum depth, the number of questions allowed along any path. A shallow tree is easy to read but may miss real patterns; a deep tree fits the training data almost perfectly and generalises worse. Cross-validation from lesson 8 is the honest way to pick the depth.

Common mistake

A split is not an explanation of cause. The tree asks about the feature that best separates the labels in this data, and that can be a stand-in for something else: ticket class and fare go together, and both say something about where on the ship a passenger slept. The second mistake is trusting a deep tree because its rules look so precise. A rule like "age between 31.5 and 32.5 and fare above some value" is usually the tree describing one or two passengers, not a pattern. Small changes in the data can also change the first split and therefore the entire tree, which makes single trees unstable.

That instability is a weakness, but it can be turned into a strength. In the next lesson we grow many different trees and let them vote.

import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier

df = pd.read_csv("data/raw/titanic.csv")
X = pd.DataFrame({"female": (df["sex"] == "female").astype(int)}).join(df[["pclass", "age", "sibsp", "parch", "fare"]])
Xtr, Xte, ytr, yte = train_test_split(X, df["survived"], test_size=0.3, stratify=df["survived"], random_state=1)
Xtr, Xte = Xtr.fillna({"age": Xtr["age"].median()}), Xte.fillna({"age": Xtr["age"].median()})  # train median for both
print(f"passengers {len(df)}, survived {df['survived'].sum()}")
for d in (1, 3, 20):
    t = DecisionTreeClassifier(max_depth=d, random_state=0).fit(Xtr, ytr)
    print(f"depth {d:2d}: train {t.score(Xtr, ytr):.3f}, test {t.score(Xte, yte):.3f}, leaves {t.get_n_leaves()}")

passengers 891, survived 342
depth  1: train 0.780, test 0.802, leaves 2
depth  3: train 0.833, test 0.802, leaves 8
depth 20: train 0.986, test 0.780, leaves 141