基尼指数计算

tree/my_tree.py

@@ -71,6 +71,15 @@ def calc_ent(x):
     return ent
 
 
+def cal_gini(x):
+    x_value_list = set([x[i] for i in range(x.shape[0])])
+    ent = 0.0
+    for x_value in x_value_list:
+        p = float(x[x == x_value].shape[0]) / x.shape[0]
+        ent += p*p
+    return ent
+
+
 def calc_ent1(x):
     """
         calculate shanno ent of x
@@ -231,6 +240,8 @@ if __name__ == '__main__':
     print(calc_ent1(Ytrain))
     print(calc_ent(Ytrain))
 
+    print("基尼指数", cal_gini(Ytrain))
+
     print(cal_ent_attr(Xtrain, Ytrain))
 
     node = fit(Xtrain, Ytrain, None, 0)