From 43b654a6a533b5404608f5b89eff3b846fdac7fe Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Micha=C5=82=20Czeka=C5=84ski?= Date: Thu, 21 May 2020 17:01:27 +0200 Subject: [PATCH] Add static method printing decision tree --- src/AI/DecisionTrees/DecisionTree.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/src/AI/DecisionTrees/DecisionTree.py b/src/AI/DecisionTrees/DecisionTree.py index 46a8fb9..0cf24b0 100644 --- a/src/AI/DecisionTrees/DecisionTree.py +++ b/src/AI/DecisionTrees/DecisionTree.py @@ -27,3 +27,14 @@ class DecisionTree(object): for branch in self.branches: if branch.label == attr.value: return branch.subtree.giveAnswer(example) + + @staticmethod + def printTree(tree, depth: int, indent: int = 20): + if isinstance(tree.root, AttributeDefinition): + print("NODE: {}".format(tree.root.name).rjust(indent * depth)) + else: + print("NODE: {}".format(str(tree.root)).rjust(indent * depth)) + + for branch in tree.branches: + print("| {}".format(str(branch.label)).rjust(indent * depth)) + DecisionTree.printTree(branch.subtree, depth + 1, indent)