From 44419b5f288de0532b4a5fb26e52f5cd5c42913d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Micha=C5=82=20Czeka=C5=84ski?= Date: Thu, 21 May 2020 17:00:59 +0200 Subject: [PATCH] Add method making decision in DecisionTree class This method takes an instance of DecisionTreeExample and returns what it "decided". --- src/AI/DecisionTrees/DecisionTree.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/src/AI/DecisionTrees/DecisionTree.py b/src/AI/DecisionTrees/DecisionTree.py index f224bda..46a8fb9 100644 --- a/src/AI/DecisionTrees/DecisionTree.py +++ b/src/AI/DecisionTrees/DecisionTree.py @@ -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)