使用基尼指数进行属性划分

tree/my_tree.py

@@ -131,6 +131,29 @@ def cal_ent_attr(Xtrain, Ytrain):
             min_mean = mean
     return min_i, min_mean, min_ent
 
+# 计算某个属性的基尼指数
+def cal_gini_attr(Xtrain, Ytrain):
+    print('sharp', Xtrain.shape)
+
+    # 对每个属性
+    min_ent = 100
+    min_i = 0
+    min_mean = 0
+    for i in range(Xtrain.shape[1]):
+        x_value_list = set([Xtrain[j][i] for j in range(Xtrain.shape[0])])
+        mean = sum(x_value_list)/len(x_value_list)
+        sum_ent = 0
+        # 二叉树
+        p = Ytrain[Xtrain[:, i] > mean]
+        sum_ent = sum_ent + cal_gini(p)*len(p)/len(Ytrain)
+        p = Ytrain[Xtrain[:, i] <= mean]
+        sum_ent = sum_ent + cal_gini(p)*len(p)/len(Ytrain)
+
+        if sum_ent < min_ent:
+            min_ent = sum_ent
+            min_i = i
+            min_mean = mean
+    return min_i, min_mean, min_ent
 
 MAX_T = 5
 
@@ -180,7 +203,7 @@ def fit(Xtrain, Ytrain, parent_node, depth):
     if depth >= MAX_T:
         return leaf_node(Ytrain)
 
-    i, mean, min_ent = cal_ent_attr(Xtrain, Ytrain)
+    i, mean, min_ent = cal_gini_attr(Xtrain, Ytrain)
     total_ent = calc_ent(Ytrain)
     print("第", i, "个属性,mean:", mean)
     # 生成节点