Add method making decision in DecisionTree class

This method takes an instance of DecisionTreeExample and returns what it "decided".
This commit is contained in:
Michał Czekański 2020-05-21 17:00:59 +02:00
parent 9d1299b162
commit 44419b5f28

View File

@ -1,7 +1,8 @@
from typing import List, Any from typing import List
from src.AI.DecisionTrees.AttributeDefinition import AttributeDefinition from src.AI.DecisionTrees.AttributeDefinition import AttributeDefinition
from src.AI.DecisionTrees.DecisionTreeBranch import DecisionTreeBranch from src.AI.DecisionTrees.DecisionTreeBranch import DecisionTreeBranch
from src.AI.DecisionTrees.DecisionTreeExample import DecisionTreeExample
class DecisionTree(object): class DecisionTree(object):
@ -16,3 +17,13 @@ class DecisionTree(object):
def addBranch(self, newBranch): def addBranch(self, newBranch):
self.branches.append(newBranch) self.branches.append(newBranch)
self.branchesNum += 1 self.branchesNum += 1
def giveAnswer(self, example: DecisionTreeExample):
if self.branchesNum == 0:
return self.root
for attr in example.attributes:
if attr.attributeDefinition.id == self.root.id:
for branch in self.branches:
if branch.label == attr.value:
return branch.subtree.giveAnswer(example)