bst_python/BST.py

class BSTNode:
    def __init__(self, value):
        self.value = value
        self.left = None
        self.right = None


class BST:
    def __init__(self):
        self.root = None
        self.size = 0

    def clear(self):
        self.root = None
        self.size = 0

    def add(self, value):
        to_add = BSTNode(value)
        if self.root is None:
            self.root = to_add
            self.size += 1
        else:
            current = self.root
            parent = None
            while current is not None:
                if value > current.value:
                    parent = current
                    current = current.right
                elif value < current.value:
                    parent = current
                    current = current.left
                else:
                    return
            if value > parent.value:
                parent.right = to_add
            else:
                parent.left = to_add
            self.size += 1

    def inOrder(self):
        ret = []
        inOrderRec(self.root, ret)
        return ret

    @classmethod
    def balance(cls, list_param):
        tree = cls()
        add_rec(0, len(list_param) - 1, list_param, tree)
        return tree

    def height(self):
        return rec_height(self.root)


def inOrderRec(top, ret):
    if top is None:
        return
    if top.left:
        inOrderRec(top.left, ret)
    ret.append(top)
    if top.right:
        inOrderRec(top.right, ret)


def rec_height(node: BSTNode):
    if node is None:
        return 0
    return 1 + max(map(rec_height, (node.left, node.right)))


def add_rec(start, end, list_param, tree: BST):
    if end - start >= 0:
        mid = (start + end) // 2
        tree.add(list_param[mid])
        add_rec(start, mid - 1, list_param, tree)
        add_rec(mid + 1, end, list_param, tree)


if __name__ == "__main__":
    test = 17
    list_param = list(range(test))
    tree = BST.balance(list_param)
    tree2 = BST()
    for i in range(test):
        tree2.add(i)
    r = list(node.value for node in tree.inOrder())
    r2 = list(node.value for node in tree2.inOrder())
    assert r == r2