Information Gain of a Split

Compute the information gain of splitting a set of labels into groups.

Decision trees such as ID3 and C4.5 pick the split that removes the most uncertainty about the label. Information gain measures exactly that: the entropy of the parent, minus the size-weighted entropy of the children.

H(S)  = -sum_c p_c * log2(p_c)
IG    = H(parent) - sum_k (|child_k| / |parent|) * H(child_k)

p_c is the fraction of labels in the set that equal class c.

Function signature

def information_gain(parent, children):

Example

Input

parent   = np.array([0, 0, 0, 1, 1, 1])
children = [np.array([0, 0, 0, 1]), np.array([1, 1])]
information_gain(parent, children)

Output

0.45914792

Reasoning

The parent is half and half, so H = 1 bit. The left child has 3 of one class and 1 of the other, so H = 0.8113. The right child is pure, so H = 0. The weighted child entropy is (4/6) * 0.8113 = 0.5409, so the gain is 1 - 0.5409 = 0.4591 bits.

Constraints and notes

• Labels are integers. The children together hold exactly the parent labels.
• Empty children are allowed and contribute nothing.
• A split that keeps the class mix in every child has a gain of 0.