Decision trees: entropy and information gain

The concept

A decision tree classifies an example by asking a series of questions about its features, following a branch for each answer until it reaches a leaf with a prediction. The tree is learned from data by choosing, at each node, the question that best separates the classes.

The ID3 algorithm measures how mixed a set of labels is with entropy, H = −Σ pᵢ log₂ pᵢ, where pᵢ is the fraction of examples in class i. A set with all labels equal has entropy 0, and a half-and-half set has entropy 1 bit. The information gain of a feature is the entropy before the split minus the weighted average entropy of the groups after it. The feature with the highest gain is chosen, and the process repeats on each group. The CART algorithm used in many libraries usually uses Gini impurity instead of entropy, with similar results.

In the example, four of eight days are 'yes' and four are 'no', so the entropy is 1.0. Splitting on outlook gives groups that are nearly pure, a gain of 0.656 bits. Splitting on windy gives almost no improvement, 0.049 bits, so outlook is chosen. A tree grown until every leaf is pure fits the training data perfectly and tends to overfit, so real implementations limit depth, require a minimum number of examples per leaf, or prune the tree.

Code example

import math

rows = [
    ("sunny", False, "no"),
    ("sunny", True, "no"),
    ("overcast", False, "yes"),
    ("rain", False, "yes"),
    ("rain", False, "yes"),
    ("rain", True, "no"),
    ("overcast", True, "yes"),
    ("sunny", False, "no"),
]


def entropy(labels):
    total = len(labels)
    return -sum((labels.count(v) / total) * math.log2(labels.count(v) / total) for v in set(labels))


def gain(rows, index):
    remainder = 0.0
    for value in set(row[index] for row in rows):
        subset = [row[-1] for row in rows if row[index] == value]
        remainder += len(subset) / len(rows) * entropy(subset)
    return entropy([row[-1] for row in rows]) - remainder


print(round(entropy([row[-1] for row in rows]), 3))
gains = {"outlook": gain(rows, 0), "windy": gain(rows, 1)}
print({name: round(value, 3) for name, value in gains.items()})
print(max(gains, key=gains.get))

1.0
{'outlook': 0.656, 'windy': 0.049}
outlook

Follow the code

Open the walkthrough to inspect the example one step at a time.

—

1.0

1.0

1.0
{'outlook': 0.656, 'windy': 0.049}

1.0
{'outlook': 0.656, 'windy': 0.049}
outlook

Behind the scenes

Cost and common mistakes

Finding the best split of a node costs O(n·d) for n examples and d categorical features. Growing a whole tree costs roughly O(n·d·depth).

• Information gain favours features with many distinct values, such as an ID column. Use gain ratio or limit such features.
• A fully grown tree memorizes the training set. Limit its depth or prune it, and check performance on held-out data.
• Splits are chosen greedily, one node at a time. The result is good but not guaranteed to be the best possible tree.

Try it yourself

What is the entropy of a set with 3 positive and 1 negative example?

Further reading

• scikit-learn user guide (external)
• An Introduction to Statistical Learning (James, Witten, Hastie, Tibshirani) (external)
• Quinlan (1986): Induction of Decision Trees (external)