完善类

tree/my_tree.py

@@ -44,6 +44,14 @@ class TreeNode(object):
             return "%s,%f, ent:%f, leaf:%d, samples:%d, value:%s, %s" % \
                    (feature_name[self.idx], self.idx_value, self.ent, self.is_leaf, self.samples, self.value, class_names[self.y])
 
+    def export(self, _feature_name, _class_names):
+        if self.idx == -1:
+            return "%s,%f, ent:%f, leaf:%d, samples:%d, value:%s, %s" % \
+               ('', self.idx_value, self.ent, self.is_leaf,self.samples, self.value, _class_names[self.y])
+        else:
+            return "%s,%f, ent:%f, leaf:%d, samples:%d, value:%s, %s" % \
+                   (_feature_name[self.idx], self.idx_value, self.ent, self.is_leaf, self.samples, self.value, _class_names[self.y])
+
 
 def read_data():
     wine = load_wine()
@@ -259,21 +267,21 @@ def distrib(Ytrain):
     return d_list
 
 
-def print_width(nodes, depth):
+def print_width(nodes, depth, _feature_name, _class_names):
     if len(nodes) == 0:
         return
 
     print("--第", depth, "层--")
     node_down = []
     for node in nodes:
-        print(node)
+        print(node.export(_feature_name, _class_names))
         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)
+    print_width(node_down, depth+1, _feature_name, _class_names)
 
 
 def predit_one(X, Y, node):
@@ -353,6 +361,12 @@ class MyDT(object):
         result = self.predict(Xtest)
         return sum(result == Ytest)/Ytest.shape[0]
 
+    @staticmethod
+    def export(tree, feature_names, class_names):
+        nodes = tree.root_node
+        print_width([nodes], 1, feature_names, class_names)
+
+
 if __name__ == '__main__':
     Xtrain, Xtest, Ytrain, Ytest = read_data()
     print(calc_ent1(Ytrain))
@@ -365,11 +379,13 @@ if __name__ == '__main__':
 
     print("信息增益率", cal_ent_attr_c45(Xtrain, Ytrain, weights))
 
-    clf = MyDT(criterion="entropy", max_depth=3,)
+    clf = MyDT(criterion="entropy", max_depth=1,)
     clf.fit(Xtrain, Ytrain, weights)
 
     # print_width([node], 1)
 
     print(clf.predict(Xtest))
 
-    print(clf.score(Xtest, Ytest))
+    print(clf.score(Xtest, Ytest))
+    print(clf.score(Xtrain, Ytrain))
+    MyDT.export(clf, feature_name, class_names)