class Test_TreeNode: def setup(self): self.t = TreeNode(key=1, val='a') @raises(ValueError) def test_value_set(self): self.t.val='?' def test_treenode(self): assert self.t.key == 1 and self.t.val == 'a' def test_treenode_root(self): assert self.t.is_root() def test_treenode_leaf(self): assert self.t.is_leaf() def test_repr(self): assert str(self.t) == f'TreeNode(key={self.t.key}, val={self.t.val})' def test_setter(self): self.t.val = 'b' assert self.t.val == 'b'
class BST(BinaryTree): def __init__(self, value): if isinstance(value, TreeNode): self.root = value elif type(value) in {set, frozenset}: lst = sorted(value) self.root = self.__init_bst(lst) elif hasattr(value, '__iter__'): lst = sorted(set(value)) self.root = self.__init_bst(lst) else: self.root = TreeNode(value) def __init_bst(self, lst): length = len(lst) assert length > 0, "Given list must have t lease on item!!" if length == 1: node = TreeNode(lst[0]) elif length == 2: node = TreeNode(lst[0]) node.set_left(TreeNode(lst[1])) else: mid_idx = len(lst) // 2 node = TreeNode(lst[mid_idx]) node.set_left(self.__init_bst(lst[0:mid_idx])) node.set_right(self.__init_bst(lst[mid_idx + 1:])) return node ############################## MAX/MIN ############################## def _get_max_node(self, start_node): # get the right-most node if start_node.get_right() == None: return start_node else: return self._get_max_node(start_node.get_right()) def get_max(self): max_node = self._get_max_node(self.root) return max_node.get_data() def _get_min_node(self, start_node): # get the left-most node if start_node.get_left() == None: return start_node else: return self._get_min_node(start_node.get_left()) def get_min(self): min_node = self._get_min_node(self.root) return min_node.get_data() ############################## SEARCH ############################## def search(self, find_val, start_node=None): if start_node is None: start_node = self.root if find_val == start_node.get_data(): return start_node elif find_val < start_node.get_data(): if start_node.get_left(): return self.search(find_val, start_node.get_left()) else: return start_node else: if start_node.get_right(): return self.search(find_val, start_node.get_right()) else: return start_node def __contains__(self, find_val): assert type(find_val) in {int, float}, "BST contains only numbers!" found_node = self.search(find_val) return found_node.get_data() == find_val ############################## INSERTION ############################## def __insert(self, value, start_node): if value == start_node.get_data(): return start_node elif value < start_node.get_data(): if start_node.get_left(): return self.__insert(value, start_node.get_left()) else: start_node.set_left(TreeNode(value)) return start_node.get_left() else: if start_node.get_right(): return self.__insert(value, start_node.get_right()) else: start_node.set_right(TreeNode(value)) return start_node.get_right() def _insert(self, value): # used mainly for inheritance assert type(value) in {int, float}, "You can insert only numbers!" return self.__insert(value, self.root) def insert(self, value): self._insert(value) ############################## REMOVAL ############################## def __remove(self, del_value, start_node): curr_value = start_node.get_data() # when del_value is found if del_value == curr_value: replacement_node = start_node.get_left() \ if not start_node.get_right() \ else self._get_min_node(start_node.get_right()) if replacement_node is None: parent = start_node.get_parent() if curr_value < parent.get_data(): parent.set_left(replacement_node) else: parent.set_right(replacement_node) return parent else: start_node.data = replacement_node.data if start_node.get_right(): self.__remove(replacement_node.data, start_node.get_right()) else: self.__remove(replacement_node.data, start_node.get_left()) return start_node # when del_value < current value elif del_value < curr_value: if start_node.get_left() is None: # raise ValueError("Couldn't find given value in the tree!!") return start_node else: return self.__remove(del_value, start_node.get_left()) # when del_value > current value else: if start_node.get_right() is None: # raise ValueError("Couldn't find given value in the tree!!") return start_node else: return self.__remove(del_value, start_node.get_right()) def _remove(self, del_value): # used mainly for inheritance assert type(del_value) in {int, float}, "BST conains numbers only!" if self.root.is_leaf() and del_value == self.root.get_data(): raise ValueError("Can't remove the only item in the tree!") return self.__remove(del_value, self.root) def remove(self, del_value): self._remove(del_value)