How to build a decision tree regressor from scratch

In the previous article we built a classifier using the decision tree to predict whether a patient has heart disease. Decision trees can also be used for regression tasks to predict a number instead of a category — a salary, a temperature, the price of a house. In this article, we will take the Hitters dataset and build a decision tree to predict a player’s salary from seasons played and hits in the previous season.

The classifier scored splits using Gini impurity reduction. For squared-error regression, we use variance reduction instead and predict the mean in each leaf. Multiplying the variance reduction by the number of rows in the parent gives the reduction in total squared error.

The split criterion and the value stored in each leaf are two sides of the same decision. A node’s impurity is the training error of the answer its leaf will give. Gini measures the cost of predicting class proportions for a pile of labels; variance measures the cost of predicting the mean for a pile of salaries. Change the leaf prediction and the appropriate error measure changes with it.

Candidate generation, recursive splitting, and the stopping rule remain unchanged. Only the impurity calculation and leaf prediction need to change.

Download the complete Python example and run python3 hitters_tree.py.

A mean-predicting leaf also limits extrapolation: every prediction stays within the range of the training targets. The later examples show why increasing tree depth cannot make this model continue a rising trend beyond the observed feature range.

The same machinery

We keep the classifier’s simplified CART procedure. Each internal node selects a binary question, scores the two resulting groups, and recurses. A leaf ends the recursion when no candidate gives a positive improvement beyond numerical tolerance.

The candidates are generated exactly as before — every feature paired with every value that feature takes in the rows at hand, one pair per question, with a numeric column asking >= and a categorical column asking == — only the target column changed type, and the predictors can still be of either kind. Nothing about that generator knows or cares what the target column contains: five rows with two columns of four distinct values each produce eight candidates whether you are predicting a disease or a salary.

What we need to change is the measure of a pile — how much the rows in it differ from one another. In the classifier that was gini, how mixed the labels in the pile are; with a numeric target it becomes variance, how far the salaries in the pile sit from their own average.

What is interesting is that the information gain formula we used to score questions does not need to change with it. This is how we scored a split in the classifier:

def info_gain(left, right, current_uncertainty):
    p = float(len(left)) / (len(left) + len(right))
    return current_uncertainty - p * gini(left) - (1 - p) * gini(right)

And this is the regressor’s, the same function with one name changed:

def info_gain(left, right, current_uncertainty):
    p = float(len(left)) / (len(left) + len(right))
    return current_uncertainty - p * variance(left) - (1 - p) * variance(right)

The function subtracts the size-weighted child variances from the parent variance. Its result is a reduction in mean squared error at this node. Here salaries are in thousands of dollars, so variance and SSE both have units of squared thousands of dollars.

The five players

The dataset is five players from the Hitters study — two numeric predictors, one numeric target:

years is seasons played in the major leagues, hits is hits in the previous season, salary is his salary for the 1987 season, in thousands of dollars. These are real rows from the real study, checked against the full file.

Measuring the spread of a dataset

To score the candidate splits, we need a measure of salary spread. Start with the sum of squared errors (SSE), then divide by the row count to obtain the variance.

Suppose you need to say which of two payrolls is the more scattered one:

Squad B comes from our real dataset — the salaries of the five players above; squad A is a hypothetical club where everyone earns about the same.

One way to do it is to put both on a graph:

One look is enough to tell that the second squad is more spread out. But the split search needs a number it can compute. The mean is the first thing to try, and it is already in the figure as the dashed line — falling in exactly the same place for both squads.

That tells us the mean cannot measure how scattered a pile is — it reports 420 either way. Instead, we can measure how far each salary sits from the mean:

We cannot use the deviations as they stand, because both sets sum to zero — a property of the mean: it is the balance point of the set, so whatever sits above it exactly cancels whatever sits below.

Squaring prevents positive and negative deviations from canceling and gives large misses more weight. Absolute error is another option, but it pairs with a median prediction rather than a mean. Differentiability is not required for the tree’s split search.

Square the deviations and add them up, and you end up with the SSE; divide that by how many there are, and you have the variance:

SSE is a total; variance is SSE per row. Duplicating every row doubles SSE without changing variance. Adding arbitrary new rows can change both. Here squad B’s variance is about 431 times squad A’s: 86,207.50 versus 200.00.

When we put that arithmetic down as a formula, variance is:

where $n$ is how many rows the pile holds, $y_i$ the salary of one of them, and $\bar{y}$ their mean — so $y_i - \bar{y}$ is one row’s deviation, the column we tabulated above.

Drop the $\frac{1}{n}$ and what is left, $\sum_{i=1}^{n}(y_i - \bar{y})^2$, is the SSE.

The formula we just built is known as the population variance, and looking it up you will find it beside a second version, the sample variance, differing only in the denominator:

Use the denominator $n$ here because we are measuring the mean squared error on the rows in this node. The $n-1$ correction serves a different purpose: estimating population variance from a random sample. It is unnecessary for this training-loss calculation.

Why this is the right impurity

Let us now understand why this particular measure is the right one to score a node with. Once the tree is built, each leaf holds the mean of its pile, and at prediction time that mean is the salary it predicts for every new player who reaches it — one number for all of them. A pile with a wide spread makes that one number badly wrong for much of what it holds, which is why the split search wants piles with as little spread as possible.

Suppose the tree never split at all — one leaf holding a whole squad, predicting the mean salary, 420.0, for every player in it.

For squad A that is a good answer: nobody there earns more than 20 away from it, so the worst the leaf can be is 20 out. For squad B it is a bad one: Robidoux earns 67.5 and Moseby 787.5, and both are told 420.0, missing by 352.5 and 367.5. Square those misses and average them and you are back at the two numbers from the table, 200.00 and 86,207.50 — the same variances, now read as the error each leaf would make. So the spread inside a node is the error that node’s answer would make — and that is the number the tree compares when it decides which question to ask, preferring the one whose two piles leave the least of it.

All of that is already in the formula, $\frac{1}{n}\sum_{i=1}^{n}(y_i - \bar{y})^2$. Here $\bar{y}$ is the 420.0 the leaf predicts for everyone who reaches it. Each $(y_i - \bar{y})$ is the error for one player: $-352.5$ for Robidoux, $+367.5$ for Moseby, and never worse than $\pm 20$ for anyone in squad A. Squaring and averaging those errors gives us 86,207.50 and 200.00. Variance is therefore the mean squared training error of a leaf that predicts the average — the error the best possible constant leaf would still make.

Classification and regression use the same structure. The leaf stores the best constant prediction, while the impurity measures the error of that prediction. For classification those are class proportions and Gini impurity; for regression they are the mean and variance. The impurity is evaluated for every candidate at every node, so changing it can move every split in the tree. The leaf statistic is not consulted when choosing splits: it determines what the finished tree predicts, not how it is shaped.

Mean or median: choosing the criterion

We chose squared error a few paragraphs ago, when the deviations needed their signs discarded and we squared them rather than taking absolute values. It is what the libraries default to, e.g. scikit-learn uses criterion="squared_error" unless you say otherwise. It is not the only option, though, and the choice reaches further than it looks: the loss we score splits with while building also defines what a leaf must store for prediction, because both are answers to the same question — which single constant minimises this loss.

• Minimise squared error → the leaf holds the mean.
• Minimise absolute error → the leaf holds the median.

Take a leaf holding [10, 12, 14, 16, 200], where the 200 is an outlier or a data-entry mistake:

The mean minimizes squared error, while the median minimizes absolute error. Raising the largest value beyond 200 would move the mean further but leave this median at 14. That resistance to an extreme value is why absolute-error leaves can be useful with outliers.

So the impurity we pick decides the constant a leaf holds, which is why the two functions below are written as a pair.

def mean(rows):
    return sum(row[-1] for row in rows) / float(len(rows))

def variance(rows):
    targets = [row[-1] for row in rows]
    m = sum(targets) / len(targets)
    return sum((t - m) ** 2 for t in targets) / len(targets)

variance scores the node; mean supplies its prediction. With absolute error, the analogous impurity would be the mean absolute deviation from the median, and the leaf would store that median.

Scoring a split

So now we know how to measure how scattered a single pile is, and we can look at how to use that measurement to score a question. Our goal is to find out how much squared error is left once the question has split the rows, or equivalently how much of it the question removed. To do that we can take either of two approaches, working in totals or in averages:

They rank the candidate questions identically, so nothing hangs on the choice. Totals are the simpler arithmetic, so we start there; the averaged form is what the code computes, and we come back to it once the criterion is in place.

A question splits the node into two piles, and each of them answers with its own mean — $\bar{y}_L$ on one side, $\bar{y}_R$ on the other. So we can apply the SSE formula to each pile separately, against its own mean. When we add those two SSEs together, we end up with the residual sum of squares (RSS) — the error the split leaves behind:

So to pick the best question out of the candidates, we compute this for each of them and take the one with the smallest RSS — and that is already enough to build a regression tree: score every candidate, keep the lowest, recurse on the two piles it made.

Looking at that formula, one might wonder whether the criterion is trying to keep the tree small, or to cluster the rows into tidy groups. It does neither: it greedily looks for feature-defined divisions that make the targets in each child easier to predict with a single leaf value — less class heterogeneity in classification, less squared variation around the mean in regression. It scores the two piles at one node, takes the winner, and recurses; what the finished tree costs as a whole is never on the table.

Scoring by what a split removes

There is a different way to use the SSE to score a question. Run the same formula on the parent pile — the rows before the split, subscript $P$ — against its own mean $\bar{y}_P$, and you get $\text{SSE}_P = \sum_{i=1}^{n_P}(y_i - \bar{y}_P)^2$ — the error that pile makes as it stands, answering every row with that one number. So instead of asking how much error a split leaves, we can ask how much it removed — that parent error minus what the two children still carry:

That is information gain’s shape, written in totals rather than in weighted averages.

The parent SSE is fixed while comparing candidates at one node. Subtracting each candidate’s residual error from that fixed value reverses the ranking: the smallest residual gives the largest gain.

The candidate that leaves the least has removed the most. Minimising the RSS and maximising the gain are doing the same job from opposite ends.

In exact arithmetic, squared-error gain cannot be negative: each child’s own mean is at least as good for its rows as the parent mean. Zero immediate gain does not imply that deeper splits could never help. Our greedy implementation stops at negligible gain; depth limits, minimum leaf sizes, and pruning are additional ways to control complexity.

The same criterion in variance notation

The gain is stated in totals, while info_gain — the classifier article’s function, with variance where it called gini — is stated in averages, weighting each child by its share of the rows:

def info_gain(left, right, current_uncertainty):
    p = float(len(left)) / (len(left) + len(right))
    return current_uncertainty - p * variance(left) - (1 - p) * variance(right)

p is the fraction of the parent’s rows that ended up in the left pile — three of five rows makes p = 0.6, leaving 1 - p = 0.4 for the right pile — and it is there because we are working in averages. A variance says nothing about how many rows produced it, so a child of one row and a child of a hundred would count equally — and if we want to keep using variances, we have to put the sizes back by hand, which is exactly what p and 1 - p do.

Written in totals no weights are needed at all:

def sse(rows):
    m = mean(rows)
    return sum((row[-1] - m) ** 2 for row in rows)

def gain_sse(rows, left, right):
    return sse(rows) - (sse(left) + sse(right))

Both score the same candidates in the same order, and the averaged form is simply what the classifier already had, Gini being an average too. The two convert into each other with one identity, since variance is SSE per row:

Substitute that for all three piles, and the gain becomes

and dividing through by $n_P$ — again a constant at this node, so again harmless — gives the form the classification article used, and the one info_gain computes:

Minimizing child SSE, maximizing SSE reduction, and maximizing variance reduction choose the same split within a fixed parent node. The values differ by a constant offset or the parent row count. Our code uses variance reduction to preserve the classifier’s scoring structure.

Scoring the root’s candidates

Let us now run that criterion on the first node of a real tree — the root, which holds all five players, before any question has been asked. Its mean is $2100/5 = 420.0$ and its variance is 86,207.50, so by $\text{SSE} = n \cdot \text{Var}$ the squared error we are trying to reduce is $\text{SSE}_P = 5 \times 86{,}207.50 = 431{,}037.5$. Every candidate below is scored against that number in the averaged form — the gain our code prints, $\text{Var}_P$ minus the two size-weighted child variances — and we read the winner back as an RSS afterwards.

Candidate questions are generated exactly as before — every feature paired with every value it takes. Two columns with four distinct values each gives eight candidates, two of which fail to divide the rows at all. They are laid out best-first here, though the code never sorts them; it just keeps a running winner:

The two skipped candidates are the smallest value in each column — years runs 2, 2, 3, 6, 7 and hits runs 41, 41, 77, 130, 149 — so every row answers yes to them. Everything goes to the True side and nothing to the False side, hence the 5 / 0: there is no split to score, and they are dropped.

The gain column is what the code reports, since info_gain is the function doing the scoring. Read the winner as total squared error instead and the criterion is easier to see. Is hits >= 77? sends Davis, Marshall and Moseby one way and the two collided players the other:

against the root’s $431{,}037.50$. One question disposes of 89% of the squared error in the dataset, and no other question on offer leaves less behind. (The gain column is the same fact per-observation: $76{,}501.0417 \times 5 = 382{,}505.21$ removed, and $431{,}037.5 - 382{,}505.21 = 48{,}532.29$.)

Two candidates tie exactly at the top — Is years >= 3? and Is hits >= 77?, both 76501.0416666667, because they cut the five players into the same two groups. The >= in find_best_split hands the win to whichever column is scanned last, exactly as it did in the classification article.

The finished tree, and what its leaves hold

The recursion stops when no candidate improves the fit beyond numerical tolerance. For these five rows, it produces:

Three leaves hold exactly one player each and reproduce their salary perfectly. The fourth holds the collided pair and answers 81.25, the mean of 67.5 and 95.0.

Run the five players back down the finished tree — the same walk a new player would take, each one following the questions to a leaf and taking the number it holds — and this is what comes back:

BillyJo Robidoux     actual    67.5   predicted   81.25
Jack Howell          actual    95.0   predicted   81.25
Alvin Davis          actual   480.0   predicted  480.00
Mike Marshall        actual   670.0   predicted  670.00
Lloyd Moseby         actual   787.5   predicted  787.50

Three training salaries are reproduced exactly, while the two players with identical predictors receive their shared mean. This describes training fit; we have not measured performance on new players.

That 81.25 clearly demonstrates why we keep the mean in a leaf rather than anything else — the smaller of the two salaries, say, or the larger. The leaf must emit one number, call it $c$, and the loss we are minimising is squared error, so the question is which $c$ makes $\sum_i (y_i - c)^2$ as small as possible. That is a smooth function of $c$, so the minimum is where its derivative vanishes:

So the mean is not one reasonable choice among several, and not a convention either: it is the only constant that satisfies this, and it is the solution to the same minimisation the split criterion is running. Impurity and leaf value come from one loss function.

That leaf is also where the tree stops improving. Robidoux and Howell have identical years and hits, so no question can ever separate them: they share a leaf at any depth, and whatever number that leaf answers, it is wrong for at least one of them. Answering 81.25 leaves $(67.5 - 81.25)^2 + (95 - 81.25)^2 = 378.125$ of squared error, and no tree reading only these two columns can drive it lower — a floor under the training error that growing deeper cannot get past.

Among the 263 players with recorded salaries, nine pairs share the same (Years, Hits) values. The largest salary difference within such a pair is $310,000, giving a minimum mean absolute error of $155,000 on that pair for any shared prediction. Additional predictors may distinguish them. This is a limitation of the recorded features, not evidence that salary is inherently unpredictable.

Inside the data: a staircase

Let us now look at the shape a finished tree draws, and then at what happens past the edge of the training data. This will help us see what a regression tree can and cannot express. A regression tree is piecewise constant: its questions divide the input space into regions, and every point in a region receives the same prediction. As a function of the features, its output is a set of flat plateaus with vertical jumps at the thresholds — no slopes anywhere, at any depth. We will refer to this property as flatness below.

To demonstrate this easily we plot a synthetic dataset below rather than Hitters: 40 rows with one feature, $x$, evenly spaced from 0 to 10, and a target that follows a smooth wave with a little noise on it.

The widget fits a regression tree to these rows. Dashed lines mark learned thresholds, and each interval takes the mean of its training targets. Increase the depth limit to see more intervals:

At depth 1 there are two plateaus and the fit is terrible. By depth 6 there are 27 and the squared error has fallen from 5,816 to 89 — that figure being the whole fit’s error, every one of the 40 points walked down the tree and scored against the mean held by the leaf it lands in, exactly the sum we have been calling the RSS, now over 27 leaves rather than two. The tree is converging on the curve, but it never bends — it approximates a smooth function by chopping it into ever-narrower constant pieces.

A tree can approximate nonlinear relationships and feature interactions without specifying them in advance. But a finite tree with constant leaves cannot represent a nonconstant straight line exactly over a continuous interval; it needs more steps for a finer approximation.

With two features, the staircase becomes a terrain

One feature gives a staircase because there is one axis to lay the steps along, with the prediction on the other. Add a second feature and both axes are spoken for, so the prediction has to go somewhere else: the tree cuts the plane into rectangles, and what was the height of a step becomes the height of a flat roof over each one.

The two panels below are the same model. On the left, the partition seen from above, where the prediction shows up as the shading and the number inside each rectangle — the picture the classification article drew for decision regions. On the right, the identical boxes raised to that number, so the height carries what the $y$ axis carried in the staircase:

Every roof is flat and every wall is vertical, which is what piecewise constant looks like when you can see it. Raise the depth and the terrain gains blocks the way the staircase gained steps — 2 regions, then 4, 8, 16 — approaching the shape of the data in flat facets, never in slopes.

Outside the data: a ceiling

In the one-feature staircase, every input above the largest learned threshold reaches the same outer leaf. With several features, increasing one feature beyond all its thresholds stops changing decisions on that feature, but the other features can still route rows to different leaves.

The mechanism behind this is not unique to regression. A classifier’s regions are just as flat, and it too answers anything beyond its training data with whatever its outermost leaf holds. For regression, the limitation is especially visible because target values have an order. If salary keeps rising outside the training range, the tree cannot follow it; it continues returning the value stored in its outermost leaf. Labels have no equivalent direction — there is no category above “disease” — so the same behaviour is less obvious in classification.

This behaviour follows directly from how prediction works: a row arriving with $x = 10^6$ answers “yes” to every threshold question on the way down, lands in the right-most leaf, and receives the mean of the training rows that landed there. There is no mechanism by which a leaf value can depend on how far past the threshold the row is.

Here is the sharpest possible demonstration — a perfectly linear relationship with no noise at all, $y = 2.5x + 3$, sampled on $x \in [0, 10]$, fitted by a depth-3 tree and by ordinary linear regression:

import numpy as np
from sklearn.linear_model import LinearRegression
from sklearn.tree import DecisionTreeRegressor

X = np.linspace(0, 10, 60).reshape(-1, 1)
y = 2.5 * X.ravel() + 3

tree = DecisionTreeRegressor(max_depth=3, random_state=0).fit(X, y)
linear = LinearRegression().fit(X, y)

Both models fit the same 60 rows, all of them inside $x \in [0, 10]$. Now ask each of them for values inside that range and then well outside it:

Drawn out to $x = 20$, with the training range ending at 10:

The linear model recovers the rule exactly and is correct at $x = 1000$. The tree returns 26.52 at $x = 12$, at $x = 50$, and at $x = 1000$ — one number, for every input beyond its experience, on data with no noise and a relationship a straight line captures with two parameters.

Inside the training range, the tree also makes approximation errors: at $x=2$ it returns 7.45 instead of 8.00. Some inputs can be predicted exactly, but eight constant plateaus cannot reproduce this line everywhere.

Worse than the flatness is the bound. A leaf holds the mean of the training targets that reached it, and a mean cannot lie outside the values it averages. So every prediction a regression tree can ever make is confined to the range of the training targets. It cannot forecast a record high or a record low, at any depth, on any data.

The line example’s largest leaf value is 26.5169, below its largest training target of 28.0. A tree can reach the target maximum if a leaf contains only that maximum. For comparison, fit on all 19 Hitters predictors, using the 184 training rows from a 70/30 split with random_state=0:

depth 2      highest possible prediction  2127.3   reached by 1 player
depth 3      highest possible prediction  2127.3   reached by 1 player
depth 5      highest possible prediction  2127.3   reached by 1 player
depth full   highest possible prediction  2127.3   reached by 1 player

In this fit, a depth-2 tree already isolates the top earner. Whether that happens depends on the predictors and the other rows, not just on how extreme the target is. The bound is unchanged: a mean-valued leaf cannot exceed the largest target it contains.

This matters most when extrapolation is part of the task:

• Trends. With other features fixed, a tree stops changing its prediction once a growing time feature passes all its thresholds. Modeling the trend separately can help.
• Random forests. An average of mean-predicting trees remains within the training target range.
• Boosted trees. Their sums can exceed that range, but a finite ensemble of constant-leaf trees is still piecewise constant. Along a fixed direction, its prediction eventually stops changing once no further split boundaries are crossed.

The limitation comes from what sits in the leaf, not from the splitting. Some tree variants fit a linear model in each leaf instead of a constant. That allows extrapolation, although predictions far beyond the data then depend heavily on the fitted slope.

From one tree to an ensemble

Turning the classifier into a regressor required only two changes: use variance to score nodes and store the mean in each leaf. The resulting model is flexible within the training range, but its piecewise-constant predictions cannot extrapolate a trend beyond it.

Regression trees are often used in ensembles. A random forest trains trees on bootstrap samples and considers random subsets of features at splits, then averages their predictions. This reduces sensitivity to any one fitted tree.

Fit them in sequence to each other’s errors instead and you get gradient boosting, where the squared error being minimised is the whole ensemble’s. Each round measures what the ensemble so far still gets wrong and fits the next tree to those residuals, so every tree runs exactly the split search from this article, only against a target made of the current mistakes rather than the raw salaries. It is also why these trees do the work even when the problem is classification: what they are fitted to is a column of real-valued gradients rather than labels.

Neither deeper trees nor additional boosting rounds guarantee zero training error: identical predictors with conflicting targets already prevent that here. Choose depth, leaf size, and ensemble size using validation performance rather than training fit alone.