Problems
Loading...
1 / 1

Compute Information Gain for a Split

Classic ML
Medium

Compute Information Gain of a binary split for classification labels. Given labels y and a boolean split_mask indicating the left/right partition, return the information gain using entropy (base 2).

Mathematical equations

Shannon Entropy (base 2):

H(Y)=kpklog2pkH(Y) = - \sum_{k} p_{k} \log_{2} p_{k}

Information Gain:

IG=H(Y)(nLNH(YL)+nRNH(YR))IG = H(Y) - \left( \frac{n_{L}}{N} H(Y_{L}) + \frac{n_{R}}{N} H(Y_{R}) \right)

where Y is the parent set of labels, YL,YR are the labels in left/right subsets after the split, and N=nL+nR.

Loading visualization...

Examples

Input: y=[0,0,1,1], mask=[T,T,F,F]

Output: 1.0 (perfect split)

Input: y=[0,0,0,1], mask=[T,F,F,T]

Output: 0.311278

Input: y=[0,0,1,2], mask=[T,F,T,F]

Output: 0.5

Hint 1

Use np.unique(y, return_counts=True) to compute label probabilities and entropy.

Hint 2

If one side is empty, return 0.0 (the split provides no information).

Requirements

  • Return a single float (information gain, base-2 entropy)
  • Handle edge cases: empty side ⇒ IG = 0.0
  • No loops over classes; NumPy only

Constraints

  • 1 ≤ N ≤ 1e6
  • Labels are integers 0..K-1 (K ≥ 2)
  • NumPy only; time limit: 200ms
Try Similar Problems