Updated DecisionTree model saving.

Author: Olcay Taner YILDIZ

Classification/Model/DecisionTree/DecisionNode.py

@@ -21,6 +21,7 @@ class DecisionNode(object):
     __class_label: str = None
     leaf: bool
     __condition: DecisionCondition
+    __classLabelsDistribution: DiscreteDistribution
     EPSILON = 0.0000000001
 
     def constructor1(self,
@@ -65,7 +66,11 @@ def constructor1(self,
         best_split_value = 0
         self.__condition = condition
         self.__data = data
-        self.__class_label = Model.getMaximum(self.__data.getClassLabels())
+        self.__classLabelsDistribution = DiscreteDistribution()
+        labels = self.__data.getClassLabels()
+        for label in labels:
+            self.__classLabelsDistribution.addItem(label)
+        self.__class_label = Model.getMaximum(labels)
         self.leaf = True
         self.children = []
         class_labels = self.__data.getDistinctClassLabels()
@@ -158,6 +163,7 @@ def constructor2(self, inputFile: TextIOWrapper):
         else:
             self.leaf = True
             self.__class_label = inputFile.readline().strip()
+            self.__classLabelsDistribution = Model.loadClassDistribution(inputFile)
 
     def __init__(self,
                  data: object,
@@ -321,7 +327,7 @@ def predict(self, instance: Instance) -> str:
 
     def predictProbabilityDistribution(self, instance: Instance) -> dict:
         if self.leaf:
-            return self.__data.classDistribution().getProbabilityDistribution()
+            return self.__classLabelsDistribution.getProbabilityDistribution()
         else:
             for node in self.children:
                 if node.__condition.satisfy(instance):

Classification/Model/DummyModel.py

@@ -22,14 +22,7 @@ def constructor1(self, trainSet: InstanceList):
 
     def constructor2(self, fileName: str):
         inputFile = open(fileName, mode='r', encoding='utf-8')
-        self.distribution = DiscreteDistribution()
-        size = int(inputFile.readline().strip())
-        for i in range(size):
-            line = inputFile.readline().strip()
-            items = line.split(" ")
-            count = int(items[1])
-            for j in range(count):
-                self.distribution.addItem(items[0])
+        self.distribution = Model.loadClassDistribution(inputFile)
         inputFile.close()
 
     def __init__(self, trainSet: object):

Classification/Model/Model.py

@@ -2,6 +2,7 @@
 from io import TextIOWrapper
 
 from DataStructure.CounterHashMap import CounterHashMap
+from Math.DiscreteDistribution import DiscreteDistribution
 from Math.Matrix import Matrix
 
 from Classification.Instance.Instance import Instance
@@ -50,6 +51,18 @@ def loadMatrix(self, inputFile: TextIOWrapper) -> Matrix:
                 matrix.setValue(j, k, float(items[k]))
         return matrix
 
+    @staticmethod
+    def loadClassDistribution(inputFile: TextIOWrapper) -> DiscreteDistribution:
+        distribution = DiscreteDistribution()
+        size = int(inputFile.readline().strip())
+        for i in range(size):
+            line = inputFile.readline().strip()
+            items = line.split(" ")
+            count = int(items[1])
+            for j in range(count):
+                distribution.addItem(items[0])
+        return distribution
+
     @staticmethod
     def getMaximum(classLabels: list) -> str:
         """