Information gain of a split
Medium12 min11 hidden tests
A churn tree is deciding whether to split customers on "used the support chat this month". You have the labels of the node before the split and of its two children. Write information_gain(parent, left, right) that returns how much the split lowers the entropy of the labels, measured in bits.
Tap a customer to move them to the other side
gain = H(parent) − weighted H(children), in bitsbefore the split · 6 customers · H = 1 bits
split on "used the support chat this month"
used chat · 4 · H = 0.811 bits
did not · 2 · H = 0 bits
Gain = 1 − 4/6 × 0.811 − 2/6 × 0 = 0.4591 bits. An empty side has entropy 0.
Gain 0.4591 bits.
Examples
The right child is pure, the left still mixed
information_gain([1, 1, 1, 0, 0, 0], [1, 1, 1, 0], [0, 0])
returns 0.459148
A split that tells you nothing gains 0 bits
information_gain(["stay", "leave", "stay", "leave"], ["stay", "leave"], ["stay", "leave"])
returns 0
Requirements
parent,leftandrightare lists of class labels (ints or strings).leftandrighttogether hold exactly the labels ofparent, and one of them may be empty.- Entropy in bits: H = -sum of p * log2(p) over the classes present, where p is a class's share. An empty list has entropy 0.0.
- Gain = H(parent) - (n_left / n) * H(left) - (n_right / n) * H(right), with n = n_left + n_right.
- Return the gain as a float.
The idea behind it
Entropy measures how uncertain a label is. In bits it is H = -sum of p * log2(p) over the classes. It is the average number of yes or no questions you need to learn a label. A pure node has entropy 0. Two classes split evenly have 1 bit, and four classes split evenly have 2 bits. A class with share 0 adds nothing, because p * log2(p) goes to 0 as p does. A split is worth the uncertainty it removes. Information gain is the parent's entropy minus the children's entropy, each child weighted by its share of the examples. A useless split gains 0, and a split that leaves pure children gains all of H(parent). The ID3 and C4.5 tree algorithms use information gain to choose splits, and scikit-learn offers it as criterion="entropy".