diff --git a/BST.py b/BST.py index 2c3a999..ddbe41f 100644 --- a/BST.py +++ b/BST.py @@ -51,15 +51,32 @@ class BST: in_order_rec(self.root, ret) return ret + def height(self): + return self.root.height() + + +class BalancedBST(BST): + def by_index(self, i): + instructions = [] + while i > 0: + instructions.append(i & 1 == 0) + i = (i - 1) // 2 + + current = self.root + while len(instructions) > 0: + go_right = instructions.pop() + if go_right: + current = current.right + else: + current = current.left + return current + @classmethod - def balance(cls, list_param): + def balanced(cls, list_param): tree = cls() add_rec(0, len(list_param) - 1, list_param, tree) return tree - def height(self): - return self.root.height() - def in_order_rec(top, ret): if top is None: @@ -82,10 +99,11 @@ def add_rec(start, end, list_param, tree: BST): if __name__ == "__main__": test = 17 list_param = list(range(test)) - tree = BST.balance(list_param) + tree = BalancedBST.balanced(list_param) tree2 = BST() for i in range(test): tree2.add(i) r = list(node.value for node in tree.in_order()) r2 = list(node.value for node in tree2.in_order()) assert r == r2 + tree.by_index(6)