A decision tree considers splitting a group of examples into smaller groups. Work out how good the split is.
Input
The first line has the labels of the parent group, separated by spaces (any words, any number of different labels). The next line has g (1 ≤ g ≤ 10), and each of the next g lines has the labels of one child group. Every parent label is in exactly one child, and no child is empty.
Output
parent H
after A
gain G
with the parent's entropy, the size-weighted entropy of the children, and the information gain, each with 3 decimal places, using log base 2.
Examples
yes yes yes yes yes yes yes yes yes no no no no no 3 yes yes no no no yes yes yes yes yes yes yes no no
parent 0.940 after 0.694 gain 0.247
a a b b 2 a a b b
parent 1.000 after 0.000 gain 1.000
Prompts inside input() are ignored when checking, so input("n: ") is fine.
Hints
Hint 1
Entropy of a group: −Σ p log₂ p over the share p of each label (math.log2).
Hint 2
The entropy after the split weights each child by its size: Σ (len(child) / len(parent)) × H(child).
Hint 3
gain = H(parent) − that weighted sum.
Solution
Show one correct solution (try it yourself first)
from math import log2
def entropy(labels):
h = 0.0
for label in set(labels):
p = labels.count(label) / len(labels)
h -= p * log2(p)
return h
parent = input().split()
g = int(input())
children = [input().split() for _ in range(g)]
after = sum(len(c) / len(parent) * entropy(c) for c in children)
print(f"parent {entropy(parent):.3f}")
print(f"after {after:.3f}")
print(f"gain {entropy(parent) - after:.3f}")