calculate information gain in decision tree algorithm

information_gain.py

# x = np.random.rand(10,2) * 10
# y = np.random.rand(10,1) * 10

def calculate_entropy(pi):    
    total = 0
    for p in pi:
        p = p / sum(pi)
        if p != 0:
            total += p * np.log2(p)
        else:
            total += 0
    total *= -1
    return total
  
def information_gain(x, y):
    total = 0
    for v in x:
        total += sum(v) / sum(y) * calculate_entropy(v)

    gain = calculate_entropy(y) - total
    return gain