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)
