扩充节点属性

tree/my_tree.py

@@ -4,23 +4,34 @@ from sklearn.datasets import load_wine
 from sklearn.model_selection import train_test_split
 import numpy as np
 
+feature_name = ['酒精', '苹果酸', '灰', '灰的碱性', '镁', '总酚', '类黄酮',
+                '非黄烷类酚类', '花青素', '颜色强度', '色调', 'od280/od315稀释葡萄酒', '脯氨酸']
+class_names=["琴酒", "雪莉", "贝尔摩德"]
+
 # 生成决策树的节点类型
 class TreeNode(object):
-    idx = ''                 # 属性
+    idx = 0                  # 属性
     idx_value = 0.0          # 属性值
     is_leaf = False
     y = 0                   # 预测值
-    next_list = []          # 分支
+    samples = 0             # 样本数
+    value = []              # 分布情况
+    left = None
+    right = None
 
-    def __init__(self, idx,idx_value, is_leaf, y, next_list):
+    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
         self.y = y
-        self.next_list = next_list
+        self.samples = samples
+        self.value = value
+        self.left = left
+        self.right = right
 
     def __str__(self):
-        return "%s,%f,%f,%d" % (self.idx, self.idx_value, self.is_leaf, self.y)
+        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])
 
 
 def read_data():
@@ -101,6 +112,7 @@ def cal_ent_attr(Xtrain, Ytrain):
             min_mean = mean
     return min_i, min_mean
 
+
 MAX_T = 5
 
 
@@ -127,7 +139,7 @@ def leaf_node(Ytrain):
         if item[1] > max_item[1]:
             max_item = item
     print('这个是叶子节点，value:', max_item[0])
-    return TreeNode('', 0, True, max_item[0], [])
+    return TreeNode(0, 0, True, max_item[0], 0, 0)
 
 
 # def create_node()
@@ -137,7 +149,7 @@ def fit(Xtrain, Ytrain, parent_node, depth):
 
     if is_end(Ytrain):
         print('这个是叶子节点')
-        return TreeNode('', 0, True, 0, [])
+        return TreeNode(0, 0, True, 0, 0, 0)
 
     if depth > MAX_T:
         return leaf_node(Ytrain)
@@ -145,22 +157,22 @@ 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, [])
+    parent_node = TreeNode(i, mean, False, 0, 0, 0)
 
     # 切分数据
     right_Ytrain = Ytrain[Xtrain[:, i] > mean]
     right_Xtrain = Xtrain[Xtrain[:, i] > mean]
-    right_Xtrain = np.delete(right_Xtrain, i, axis=1)
+    # right_Xtrain = np.delete(right_Xtrain, i, axis=1) # 这个属性还可以再被切分
 
     right_node = fit(right_Xtrain, right_Ytrain, parent_node, depth+1)
 
     left_Ytrain = Ytrain[Xtrain[:, i] <= mean]
     left_Xtrain = Xtrain[Xtrain[:, i] <= mean]
-    left_Xtrain = np.delete(left_Xtrain, i, axis=1)
+    # left_Xtrain = np.delete(left_Xtrain, i, axis=1)
     left_node = fit(left_Xtrain, left_Ytrain, parent_node, depth + 1)
 
-    parent_node.next_list.append(left_node)
-    parent_node.next_list.append(right_node)
+    parent_node.left = left_node
+    parent_node.right = right_node
     return parent_node
 
 
@@ -172,11 +184,15 @@ def print_width(nodes, depth):
     node_down = []
     for node in nodes:
         print(node)
-        if len(node.next_list) > 0:
-            node_down.extend(node.next_list)
+        if node.left is not None:
+            node_down.append(node.left)
+        if node.right is not None:
+            node_down.append(node.right)
+
 
     print_width(node_down, depth+1)
 
+
 if __name__ == '__main__':
     Xtrain, Xtest, Ytrain, Ytest = read_data()
     print(calc_ent1(Ytrain))