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.DecisionTreeBranch import DecisionTreeBranch
from src.AI.DecisionTrees.DecisionTreeExample import DecisionTreeExample
class DecisionTree(object):
@ -16,3 +17,13 @@ class DecisionTree(object):
def addBranch(self, newBranch):
self.branches.append(newBranch)
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)