Information Gain 计算 python实现

from math import log

def I(*args):
    total = sum(args) + 0.0
    result = 0.0
    for i in args:
        if i == 0:
            result += 0
        else:
            result += i / total * log( i / total, 2)
    return -result

#num表示分类的个数
def E(num, *args):
    if len(args) % num != 0:
        print "Error len(args)"
    result = 0.0
    total = sum(args)
    for x in xrange(len(args) / num):
        k = x * num
        total_up = 0.0 + sum(args[k:k + num])
        result += total_up / total * I(*args[k:k + num])
    return result

def Gain(i, e):
    return i - e

#取个别名
i = I
e = E
g = Gain

if __name__ == "__main__":
    #for example
    print i(9, 5)
    print e(2, 2, 3, 4, 0, 3, 2)