重构 类

tree/my_tree.py

@@ -47,7 +47,7 @@ class TreeNode(object):
 
 def read_data():
     wine = load_wine()
-    print(wine.data.shape)  # 178*13
+    print("数组结构", wine.data.shape)  # 178*13
     print(wine.target)
     print(wine.feature_names)
     print(wine.target_names)
@@ -148,7 +148,7 @@ def cal_max_ent_attr_c45(Xtrain, Ytrain, weights=None):
     weights = weights / sum(weights)
     h = calc_ent(Ytrain, weights)
     p = 0
-    for k in range(len(Xtrain) - 1):
+    for k in range(0, len(Xtrain) - 1, 3):
         left = Xtrain[:k + 1]
         right = Xtrain[k + 1:]
 
@@ -259,40 +259,6 @@ def distrib(Ytrain):
     return d_list
 
 
-def fit(Xtrain, Ytrain, parent_node, depth, weights):
-
-    if is_end(Ytrain):
-        # print('这个是叶子节点')
-        return leaf_node(Ytrain, weights)
-
-    if depth >= MAX_T:
-        return leaf_node(Ytrain, weights)
-
-    i, mean, min_ent = cal_ent_attr(Xtrain, Ytrain, weights)
-    total_ent = 0 # calc_ent(Ytrain)
-    # print("第", i, "个属性,mean:", mean)
-    # 生成节点
-    parent_node = TreeNode(i, mean, total_ent - min_ent, False, -2, len(Ytrain), distrib(Ytrain))
-
-    # 切分数据
-    right_position = Xtrain[:, i] > mean
-    right_Ytrain = Ytrain[right_position]
-    right_Xtrain = Xtrain[right_position]
-    # right_Xtrain = np.delete(right_Xtrain, i, axis=1) # 这个属性还可以再被切分
-
-    right_node = fit(right_Xtrain, right_Ytrain, parent_node, depth+1, weights[right_position])
-
-    left_position = Xtrain[:, i] <= mean
-    left_Ytrain = Ytrain[left_position]
-    left_Xtrain = Xtrain[left_position]
-    # left_Xtrain = np.delete(left_Xtrain, i, axis=1)
-    left_node = fit(left_Xtrain, left_Ytrain, parent_node, depth + 1, weights[left_position])
-
-    parent_node.left = left_node
-    parent_node.right = right_node
-    return parent_node
-
-
 def print_width(nodes, depth):
     if len(nodes) == 0:
         return
@@ -313,8 +279,8 @@ def print_width(nodes, depth):
 def predit_one(X, Y, node):
     if node.is_leaf:
         # print(class_names[node.y], class_names[Y])
-        if node.y == 0:
-            return -1
+        # if node.y == 0:
+        #     return -1
         return node.y
     else:
         if X[node.idx] <= node.idx_value:
@@ -323,12 +289,69 @@ def predit_one(X, Y, node):
             return predit_one(X, Y, node.right)
 
 
-def predict(Xtest, Ytest, node):
-    result = []
-    for i in range(Xtest.shape[0]):
-        result.append(predit_one(Xtest[i], None, node))
-    return np.array(result)
+class MyDT(object):
+
+    criterion = None
+    max_depth = None
 
+    root_node = None
+
+    def __init__(self, criterion, max_depth):
+        self.criterion = criterion
+        self.max_depth = max_depth
+
+    def fit(self, Xtrain, Ytrain, sample_weight=None):
+        if sample_weight is None:
+            sample_weight = np.ones(Ytrain.shape[0]) / Ytrain.shape[0]
+        self.root_node = self.do_fit(Xtrain, Ytrain, 0, sample_weight)
+
+    def do_fit(self, Xtrain, Ytrain, depth, weights):
+
+        if is_end(Ytrain):
+            # print('这个是叶子节点')
+            return leaf_node(Ytrain, weights)
+
+        if depth >= self.max_depth:
+            return leaf_node(Ytrain, weights)
+
+        if self.criterion == 'entropy':
+            i, mean, min_ent = cal_ent_attr(Xtrain, Ytrain, weights)
+        elif self.criterion == 'C4.5':
+            i, mean, min_ent = cal_ent_attr_c45(Xtrain, Ytrain, weights)
+        else:
+            i, mean, min_ent = cal_gini_attr(Xtrain, Ytrain, weights)
+        total_ent = 0  # calc_ent(Ytrain)
+        # print("第", i, "个属性,mean:", mean)
+        # 生成节点
+        parent_node = TreeNode(i, mean, total_ent - min_ent, False, None, len(Ytrain), distrib(Ytrain))
+
+        # 切分数据
+        right_position = Xtrain[:, i] > mean
+        right_Ytrain = Ytrain[right_position]
+        right_Xtrain = Xtrain[right_position]
+        # right_Xtrain = np.delete(right_Xtrain, i, axis=1) # 这个属性还可以再被切分
+
+        right_node = self.do_fit(right_Xtrain, right_Ytrain, depth + 1, weights[right_position])
+
+        left_position = Xtrain[:, i] <= mean
+        left_Ytrain = Ytrain[left_position]
+        left_Xtrain = Xtrain[left_position]
+        # left_Xtrain = np.delete(left_Xtrain, i, axis=1)
+        left_node = self.do_fit(left_Xtrain, left_Ytrain, depth + 1, weights[left_position])
+
+        parent_node.left = left_node
+        parent_node.right = right_node
+        return parent_node
+
+    def predict(self, Xtest):
+        result = []
+        for i in range(Xtest.shape[0]):
+            result.append(predit_one(Xtest[i], None, self.root_node))
+        return np.array(result)
+
+    def score(self, Xtest, Ytest):
+        result = self.predict(Xtest)
+        return sum(result == Ytest)/Ytest.shape[0]
 
 if __name__ == '__main__':
     Xtrain, Xtest, Ytrain, Ytest = read_data()
@@ -342,7 +365,11 @@ if __name__ == '__main__':
 
     print("信息增益率", cal_ent_attr_c45(Xtrain, Ytrain, weights))
 
-    node = fit(Xtrain, Ytrain, None, 0, weights)
-    print_width([node], 1)
+    clf = MyDT(criterion="entropy", max_depth=3,)
+    clf.fit(Xtrain, Ytrain, weights)
+
+    # print_width([node], 1)
+
+    print(clf.predict(Xtest))
 
-    print(predict(Xtest, Ytest, node))
+    print(clf.score(Xtest, Ytest))