Binary Tree

very gay :c
mail@pastecode.io avatar
unknown
python
2 years ago
1.5 kB
5
Indexable
Never
class BinaryNode:
    def __init__(self, key: int, left = None, right = None) -> None:
        self.key = key
        self.left = left
        self.right = right

class BinaryTree:
    def __init__(self, root: int = None) -> None:
        self.data = {}
        if root is not None:
            self.root = BinaryNode(root)
            self.data[str(root)] = self.root
    
    def __repr__(self) -> str:
        return str(self.traverse(self.root))
        
    def traverse(self, node, reachable = None) -> list:
        if not node or str(node.key) not in self.data: return
        if reachable is None: reachable = []

        self.traverse(node.left, reachable)
        reachable.append(node.key)
        self.traverse(node.right, reachable)
        return reachable
    
    def add(self, key: int, left = None, right = None) -> bool:
        if str(key) not in self.data: raise Exception("Error: Key couldn't be located in the BinaryTree")

        if left is not None:
            if str(left) not in self.data: 
                self.data[str(left)] = BinaryNode(left)
            self.data[str(key)].left = self.data[str(left)]
        else:
            self.data[str(key)].left = None
        
        if right is not None:
            if str(right) not in self.data:
                self.data[str(right)] = BinaryNode(right)
            self.data[str(key)].right = self.data[str(right)]
        else:
            self.data[str(key)].right = None