simple tree using dictionaries

simple_tree_using_dictionaries.py

class DecisionTree(object):
    def __init__(self, max_depth=2):
        self.max_depth = max_depth
        self.tree = {}
        self.sum = 0
        self.data = []
        
    def __str__(self):
        return f"data: {self.data}"
    
    def learn(self, data=np.random.rand(), depth=0):
        self.tree = self._learn(self.tree, data=np.random.rand(), depth=0)
    
    def _learn(self, tree, data=np.random.rand(), depth=0):
        if depth <= self.max_depth:
            depth+=1
            tree['data'] = data
            tree['left'] = self._learn(DecisionTree(self.max_depth).tree, data=np.random.randint(0,10), depth=depth)
            tree['right'] = self._learn(DecisionTree(self.max_depth).tree, data=np.random.randint(0,50), depth=depth)
        return tree

t = DecisionTree(max_depth=5)
t.learn()
pp.pprint(t.tree)