Information Gain

Problem

Compute the information gain (base 2) of a split of labels into left_idx and right_idx.

IG = H(parent) - (n_left/n) * H(left) - (n_right/n) * H(right), where H is Shannon entropy and an empty child has entropy 0. Empty labels return 0.0.

Hint

Weight each child's entropy by its share of the samples and subtract from the parent entropy.

case 1

in: [[0,0,1,1],[0,1],[2,3]]

expected: 1

case 2

in: [[0,0,0,1],[0,1,2],[3]]

expected: 0.8112781244591328

case 3

in: [[0,1,0,1],[0,2],[1,3]]

expected: 1

case 4

in: [[0,1,2,3],[0,1],[2,3]]

expected: 1

Starter codePython

import math
def information_gain(labels, left_idx, right_idx):
    # Your code here
    pass

4 test cases · run your code against them in the editor