Information Gain of a Split

Problem

Compute the information gain in bits of a split defined by groups, where each group is a list of indices into labels. The gain is the parent entropy minus the size-weighted average of the child entropies. Return 0.0 when groups is empty or labels is empty.

Hint

Weight each child entropy by that child's share of the samples.

Test cases

case 1

in: [["a","a","b","b"],[[0,1],[2,3]]]

expected: 1

case 2

in: [["a","a","b","b","c"],[[0,1,2],[3,4]]]

expected: 0.5709505944546684

case 3

in: [["a","b"],[]]

expected: 0

case 4

in: [["a","a","a"],[[0,1,2]]]

expected: 0

case 5

in: [["a","b","a","c"],[[0,2],[1,3]]]

expected: 1

Starter codePython

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

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