diff --git a/bst.py b/bst.py index bb6d302..08b352b 100755 --- a/bst.py +++ b/bst.py @@ -13,7 +13,7 @@ def __init__(self, value=None): self._depth = 0 if value is not None: - self.tree[value] = {'depth': 1} + self.tree[value] = {} self.top = value self._size += 1 self._depth += 1 @@ -21,10 +21,11 @@ def __init__(self, value=None): def left(self, current): return self.tree[current].get('left') - def right(self, current): return self.tree[current].get('right') + def parent(self, current): + return self.tree[current].get('parent') def insert(self, value): """Insert a node with value in order. @@ -32,7 +33,7 @@ def insert(self, value): In value is already present, it is ignored.""" current = self.top if current is None: - self.tree[value] = {'depth': 1} + self.tree[value] = {} self.top = value self._size += 1 self._depth += 1 @@ -50,14 +51,111 @@ def insert(self, value): child = 'left' if traverse is None: #actual insert - self.tree[value] = {'depth': depth} + self.tree[value] = {} + self.tree[current][child] = value + self.tree[value]['parent'] = current self._size += 1 if depth > self._depth: self._depth = depth return current = traverse + def delete(self, val): + if self.contains(val): + left_child = self.left(val) + right_child = self.right(val) + parent = self.parent(val) + if left_child is None and right_child is None: + # Delete a node with not children + if parent is not None: + # For a node with a parent... + if self.left(parent) == val: + del self.tree[parent]['left'] + else: + del self.tree[parent]['right'] + else: + # For deleting a node that is the root with no children + self.top = None + del self.tree[val] + elif left_child is not None: + biggest_child = self._rightmost(left_child) + self._swap_nodes(val, biggest_child) + self.delete(val) + else: + if parent is not None: + if self.left(parent) == val: + self.tree[parent]['left'] = right_child + else: + self.tree[parent]['right'] = right_child + self.tree[right_child]['parent'] = parent + else: + self.top = right_child + del self.tree[val] + + + def _swap_nodes(self, current, target): + """ + Swap two nodes. + + Additionally, update all information on related nodes, as in, children + and parent nodes. + """ + cur_parent = self.parent(current) + targ_parent = self.parent(target) + + children = [] + + children.append(self.left(current)) + children.append(self.right(current)) + children.append(self.left(target)) + children.append(self.right(target)) + + for child in children: + # If a node children, update it's childrens' information + # for 'swapping' + if child and self.parent(child) == current: + self.tree[child]['parent'] = target + elif child: + self.tree[child]['parent'] = current + + # Swapping information in the parent nodes relative to the nodes being 'swapped' + if current != self.top: + if current == self.tree.get(cur_parent).get('left'): + self.tree[cur_parent]['left'] = target + else: + self.tree[cur_parent]['right'] = target + + if target == self.tree.get(targ_parent).get('left'): + self.tree[targ_parent]['left'] = current + else: + self.tree[targ_parent]['right'] = current + + # Swap current and target's left, right, and parent information + self.tree[current], self.tree[target] = self.tree[target], self.tree[current] + + if current == self.top: + self.top = target + + + def _rightmost(self, start): + """Returns None if start is empty tree""" + while start is not None: + if self.right(start) is None: + return start + start = self.right(start) + + + + # left_of_parent = self.left(self.parent(current)) + # right_of_parent = self.right(self.parent(current)) + + # if current == val: + # if current == left_of_parent: + # del self.tree[self.parent(current)]['left'] + # else: + # del self.tree[self.parent(current)]['right'] + def balance(self): """Returns the balance of the tree: @@ -70,11 +168,11 @@ def balance(self): right_deep = 0 for node, v in self.tree.items(): if self.top > node: - if left_deep < v['depth']: - left_deep = v['depth'] + if left_deep < self.node_depth(node): + left_deep = self.node_depth(node) elif self.top < node: - if right_deep < v['depth']: - right_deep = v['depth'] + if right_deep < self.node_depth(node): + right_deep = self.node_depth(node) return right_deep - left_deep def contains(self, value): @@ -89,6 +187,14 @@ def depth(self): """Returns the depth of the tree.""" return self._depth + def node_depth(self, current): + """Returns the depth of a node in the tree.""" + depth = 0 + while current is not None: + current = self.parent(current) + depth +=1 + return depth + def get_dot(self): """return the tree with root 'self' as a dot graph for visualization""" return "digraph G{\n%s}" % ("" if not self.tree else ( @@ -103,6 +209,10 @@ def _get_dot(self, current): """recursively prepare a dot graph entry for this node.""" left = self.tree[current].get('left') right = self.tree[current].get('right') + parent = self.tree[current].get('parent') + if parent is not None: + yield "\t%s -> %s;" % (current, parent) + if left is not None: yield "\t%s -> %s;" % (current, left) for i in self._get_dot(left): @@ -120,6 +230,7 @@ def _get_dot(self, current): yield "\tnull%s [shape=point];" % r yield "\t%s -> null%s;" % (current, r) + def in_order(self, current='start'): """ Generator that traverses the binary tree in order. @@ -176,6 +287,8 @@ def main(): inserts = [7, 4, 11, 2, 9, 6, 12, 5, 13, 0, 10, 8, 3, 1] for i in inserts: tree.insert(i) + + tree._swap_nodes(6, 4) print tree.tree # for num in enumerate(tree.pre_order()): # print num diff --git a/test_bst.py b/test_bst.py index d1d546c..ee0ba17 100644 --- a/test_bst.py +++ b/test_bst.py @@ -7,34 +7,31 @@ def test_Bst(): tree = bst.Bst() assert tree.tree == {} tree = bst.Bst(5) - assert tree.tree == {5: {'depth': 1}} + assert tree.tree == {5: {}} def test_insert(empty_tree): """Test insert into a tree""" tree = empty_tree tree.insert(20) - assert tree.tree == {20: {'depth': 1}} + assert tree.tree == {20: {}} tree.insert(30) - expect = {20: {'depth': 1, - 'right': 30}, - 30: {'depth': 2}} + expect = {20: {'right': 30}, + 30: {'parent': 20}} assert tree.tree == expect - tree.insert(40) - expect = {20: {'depth': 1, - 'right': 30}, - 30: {'depth': 2, - 'right': 40}, - 40: {'depth': 3}} + tree.insert(40) + expect = {20: {'right': 30}, + 30: {'right': 40, + 'parent': 20}, + 40: {'parent': 30}} assert tree.tree == expect tree.insert(10) - expect = {20: {'depth': 1, - 'left': 10, + expect = {20: {'left': 10, 'right': 30}, - 30: {'depth': 2, - 'right': 40}, - 40: {'depth': 3}, - 10: {'depth': 2}} + 30: {'right': 40, + 'parent': 20}, + 40: {'parent': 30}, + 10: {'parent': 20}} assert tree.tree == expect @@ -119,6 +116,86 @@ def test_breadth_first_order(filled_tree): assert expected_order[place] == item +def test_swap(filled_tree_2): + tree = filled_tree_2 + tree._swap_nodes(7, 5) + assert tree.top == 5 + + assert tree.parent(5) is None + assert tree.left(5) == 4 + assert tree.right(5) == 11 + + assert tree.parent(7) == 6 + assert tree.left(7) is None + assert tree.right(7) is None + +def test_swap_parentchild(filled_tree_2): + tree = filled_tree_2 + tree._swap_nodes(7, 4) + assert tree.top == 4 + + assert tree.parent(4) is None + assert tree.left(4) == 7 + assert tree.right(4) == 11 + + assert tree.parent(7) == 4 + assert tree.left(7) is 2 + assert tree.right(7) is 6 + +def test_deletion_easy_case(filled_tree_2): + tree = filled_tree_2 + assert tree.left(6) == 5 + tree.delete(5) + assert tree.left(6) is None + assert tree.contains(5) is False + + +def test_deletion_one_left_child_lesser(filled_tree_2): + tree = filled_tree_2 + assert tree.parent(6) == 4 + assert tree.left(6) == 5 + assert tree.right(6) is None + tree.delete(6) + assert tree.left(4) == 2 + assert tree.right(4) == 5 + + assert tree.parent(5) == 4 + assert tree.left(5) is None + assert tree.right(5) is None + + assert tree.contains(6) is False + + +def test_deletion_one_right_child_lesser(filled_tree_2): + tree = filled_tree_2 + assert tree.parent(0) == 2 + assert tree.left(0) is None + assert tree.right(0) == 1 + tree.delete(0) + + assert tree.left(2) == 1 + assert tree.right(2) == 3 + + assert tree.parent(1) == 2 + assert tree.left(1) is None + assert tree.right(1) is None + + assert tree.contains(0) is False + + +def test_node_depth(filled_tree_2): + assert filled_tree_2.node_depth(7) == 1 + assert filled_tree_2.node_depth(4) == 2 + assert filled_tree_2.node_depth(1) == 5 + assert filled_tree_2.node_depth(13) == 4 + +def test_rightmost(filled_tree_2): + assert filled_tree_2._rightmost(0) == 1 + assert filled_tree_2._rightmost(4) == 6 + assert filled_tree_2._rightmost(7) == 13 + assert filled_tree_2._rightmost(6) == 6 + + @pytest.fixture(scope='function') def filled_tree(): """Upside down V Shaped Tree"""