binary-tree.py

class BinaryTreeNode(object):
    value = None
    left = None
    right = None

    def __init__(self, value):
        self.value = value

    def generate_dot_diagram(self, graph=None):
        if not graph:
            graph = Digraph(comment=None)

        graph.node(str(self.value), str(self.value))
        if self.left:
            self.left.generate_dot_diagram(graph)
            graph.edge(str(self.value), str(self.left.value))
        if self.right:
            self.right.generate_dot_diagram(graph)
            graph.edge(str(self.value), str(self.right.value))

        return graph


class BinaryTree(object):
    head = None
    count = 0

    def add(self, value):
        if not self.head:
            self.head = BinaryTreeNode(value)
        else:
            self.add_to(self.head, value)
        self.count += 1

    def add_to(self, node, value):
        # case 1: value is less than the current node value
        if value < node.value:
            if not node.left:
                node.left = BinaryTreeNode(value)
            else:
                self.add_to(node.left, value)

        # case 2: value is greater than or equal to current node value
        else:
            if not node.right:
                node.right = BinaryTreeNode(value)
            else:
                self.add_to(node.right, value)

        self.count += 1

    def contains(self, value):
        current, parent = self.find_with_parent(value)
        return True if current else False

    def find_with_parent(self, value):
        current = self.head
        parent = None

        while current:
            if current.value == value:
                break
            # if value is less than current, go left
            elif value < current.value:
                parent = current
                current = current.left
            # otherwise if value is greater than current, go right
            else:
                parent = current
                current = current.right

        return current, parent

    def remove(self, value):
        current, parent = self.find_with_parent(value)

        if not current:
            return None

        self.count -= 1

        # case 1: if current has no right child, then left child replaces current
        if not current.right:
            if not parent:
                self.head = current.left
            else:
                # if parent value is greater than current value, make the current left a left child of parent
                if parent.value > current.value:
                    parent.left = current.left
                # otherwise if the parent value is less than the current value, make the current left child a right child of parent
                elif parent.value < current.value:
                    parent.right = current.left

        # case 2: if the current's right child has no left child, then current's right child replaces current
        elif not current.right.left:
            current.right.left = current.left
            if not parent:
                self.head = current.right
            else:
                # if parent's value is greater than current value make the current right child a left child of parent
                if parent.value > current.value:
                    parent.left = current.right
                # otherwise if the parent value is less than the current value make the current right child a right child of parent
                elif parent.value < current.value:
                    parent.right = current.right

        # case 3: if the current's right child has a left child, replace current with current's right child's left-most child
        else:
            left_most = current.right.left
            left_most_parent = current.right

            while left_most.left:
                left_most_parent = leftmost
                left_most = left_most.left

            # the parent's left subtree becomes the leftmost's right subtree
            left_most_parent.left = left_most.right

            # assign leftmost's left and right to current's left and right children
            left_most.left = current.left
            left_most.right = current.right

            if not parent:
                self.head = left_most
            else:
                # if parent value is greater than current value, make leftmost the parent's left child
                if parent.value > current.value:
                    parent.left = left_most
                # if parent's value is less than current value, make leftmost the parent's right child
                elif parent.value < current.value:
                    parent.right = left_most

        return True

tree = BinaryTree()
tree.add(4)
tree.add(2)
tree.add(3)
tree.add(1)
tree.add(6)
tree.add(5)
tree.add(9)
tree.add(7)
tree.add(8)
tree.add(10)

# ---------------------------------------------------------------------------
# OPTIONAL: generate graph diagram in DOT graph description diagram
# requires graphviz module, which can be installed with `pip install graphviz`
# `.dot` files may then be opened with the open source Graphviz application
# --------------------------------------------------------------------------
# from graphviz import Digraph
# dot = tree.head.generate_dot_diagram()
# f = open('tree_copy.dot', 'w')
# f.write(dot.source)
# f.close()

tree.remove(6)

# dot = tree.head.generate_dot_diagram()
# f = open('tree_copy_remove.dot', 'w')
# f.write(dot.source)
# f.close()