信息熵、条件熵、信息增益、基尼指数的python实现

手动编写计算信息熵、条件熵、信息增益、基尼指数的函数

data = pd.read_excel('customer.xlsx', engine='openpyxl')

# 信息熵
# data -> 数据框 D -> 指定变量
# 返回指定变量对应的信息熵
def H(data, D):
    pk = pd.value_counts(data[D])/len(data[D])
    return sum(pk*np.log2(pk)*(-1))

# 条件熵
# data -> 数据框 D -> 指定变量 A -> 条件变量
# 返回指定变量对应的条件熵
def H_(data, D, A):
    HA = data.groupby(A).apply(lambda x:H(x, D))
    Pk = pd.value_counts(data[A])/len(data[A])
    # print(Pk)
    return sum(HA * Pk)

#信息增益
def informincrease(data,D,A):
    return H(data,D) - H_(data, D, A)

def Gini(data, D, A):
    Pk = data.groupby(A).apply(lambda x: pd.value_counts(x[D])/len(x[D]))
    Pa = pd.value_counts(data[A])/len(data[A])
    return sum(Pa * (1-Pk.groupby(A).apply(lambda x: sum(x**2))))
print(Gini(data, 'Buy', 'Stu'))