打印属性

tree/my_tree.py

@@ -19,7 +19,7 @@ class TreeNode(object):
     left = None
     right = None
 
-    def __init__(self, idx,idx_value, is_leaf, y, samples, value, left=None, right=None):
+    def __init__(self, idx, idx_value, is_leaf, y, samples, value, left=None, right=None):
         self.idx = idx
         self.idx_value = idx_value
         self.is_leaf = is_leaf
@@ -29,6 +29,10 @@ class TreeNode(object):
         self.left = left
         self.right = right
 
+        if self.y == -1:
+            self.y = np.where(value == np.max(value))[0][0]
+            print(self.y, self.value)
+
     def __str__(self):
         return "%s,%f, leaf:%d, samples:%d, value:%s, %s" % \
                (feature_name[self.idx], self.idx_value, self.is_leaf,self.samples, self.value, class_names[self.y])
@@ -139,17 +143,24 @@ def leaf_node(Ytrain):
         if item[1] > max_item[1]:
             max_item = item
     print('这个是叶子节点，value:', max_item[0])
-    return TreeNode(0, 0, True, max_item[0], 0, 0)
+    return TreeNode(0, 0, True, max_item[0], len(Ytrain), distrib(Ytrain))
+
 
+def distrib(Ytrain):
+    x_value_list = set([Ytrain[i] for i in range(Ytrain.shape[0])])
+    ent = 0.0
+    d_list = np.zeros(3, dtype=int)
+    for x_value in x_value_list:
+        d_list[x_value] = len([1 for i in Ytrain == x_value if i])
 
-# def create_node()
+    return d_list
 
 
 def fit(Xtrain, Ytrain, parent_node, depth):
 
     if is_end(Ytrain):
         print('这个是叶子节点')
-        return TreeNode(0, 0, True, 0, 0, 0)
+        return TreeNode(0, 0, True, -1, len(Ytrain), distrib(Ytrain))
 
     if depth > MAX_T:
         return leaf_node(Ytrain)
@@ -157,7 +168,7 @@ def fit(Xtrain, Ytrain, parent_node, depth):
     i, mean = cal_ent_attr(Xtrain, Ytrain)
     print("第", i, "个属性,mean:", mean)
     # 生成节点
-    parent_node = TreeNode(i, mean, False, 0, 0, 0)
+    parent_node = TreeNode(i, mean, False, -1, len(Ytrain), distrib(Ytrain))
 
     # 切分数据
     right_Ytrain = Ytrain[Xtrain[:, i] > mean]