2
\$\begingroup\$

Just a general code review. I'm looking for my first developer job so anything constructive is awesome. Please, and thank you very much.

I've written this implementation as an exercise and to add to my code base. I've tried to keep most nuts and bolts hidden away in the Node class.

"""AVLTree/Balanced Binary Search Tree Data structure.

Thanks for checking out my implementation of a BBST. Feel free to use it and/or change it to better suit your needs.
Any feedback is greatly appreciated.

Tree designed for data types supporting <, =, >.

The following methods are implemented:

    Create a new tree:
        tree = AVLTree()

    View tree structure (may need some adjustment for screen size):
        print(tree)

    Determine height of tree:
        tree.height(print_result=True)
        Returns int
        print_result is optional, defaults to False

    Insert data into tree:
        tree.insert(data)
        If data already exists in tree, user is alerted

    Search for data in tree:
        tree.search(data, print_results=True)
        returns Boolean
        print_results is optional, defaults to False

    Delete data from tree:
        tree.delete(data)
        If data does not exist in tree, user is alerted

    Clear all data from tree:
        tree.clear_tree()
"""


class Node(object):

    """Node object for AVLTree.

    Node class is wrapped by AVLTree class. All user methods are exposed there.
    Methods for printing, searching, inserting, deleting and determining height of AVLTree provided.

    """

    def __init__(self, data):
        """Instantiates Node object for AVLTree.

        Data assumed to be compatible with <, =, > operators.
        Parent, left and right are pointers to parent and child nodes.
        Tallness initialized at 1, adjusted with insert and delete to represent height of node. Null nodes have
        height of 0.

        :param data: Values compatible with <, =, > operators.
        """
        self.data = data
        self.parent = None
        self.left = None
        self.right = None
        self.tallness = 1

    def print_tree(self, cur_node, order=None):
        """Prints tree in specified order to stdout.

        :param cur_node: Root node from Tree.get_root.
        :param order: Keyword arg for order of tree traversal.
        """
        if order == 'pre-order':
            print(*self._pre_order(cur_node, []))
        elif order == 'in-order':
            print(*self._in_order(cur_node, []))
        elif order == 'post-order':
            print(*self._post_order(cur_node, []))

    def _pre_order(self, cur_node, output):
        """Prints tree: Root -> Left node -> Right node.

        :param cur_node: Root of tree.
        :param output: Empty list.
        :return: List of node values.
        """
        if cur_node:
            output += [cur_node.data]
            self._pre_order(cur_node.left, output)
            self._pre_order(cur_node.right, output)
        return output

    def _in_order(self, cur_node, output):
        """Prints tree: Left node -> Root -> Right node.

        :param cur_node: Node.Root of tree.
        :param output: Empty list.
        :return: List of node values.
        """
        if cur_node:
            self._in_order(cur_node.left, output)
            output += [cur_node.data]
            self._in_order(cur_node.right, output)
        return output

    def _post_order(self, cur_node, output):
        """Prints tree: Left node -> Right node -> Root.

        :param cur_node: Root of tree.
        :param output: Empty list.
        :return: List of node values.
        """
        if cur_node:
            self._post_order(cur_node.left, output)
            self._post_order(cur_node.right, output)
            output += [cur_node.data]
        return output

    def height(self, cur_node, cur_height):
        """Calculates and returns total height of tree.

        :param cur_node: Root of tree.
        :param cur_height:  0.
        :return: Height of tree.
        """
        if not cur_node:
            return cur_height
        left_height = self.height(cur_node.left, cur_height + 1)
        right_height = self.height(cur_node.right, cur_height + 1)
        return max(left_height, right_height)

    def search(self, cur_node, data):
        """Searches tree for data, returns bool.

        :param cur_node: Root of tree.
        :param data: Data to be found in tree.
        :return: bool. True if data in tree, else, False.
        """
        if data == cur_node.data:
            return True
        elif data < cur_node.data and cur_node.left:
            return self.search(cur_node.left, data)
        elif data > cur_node.data and cur_node.right:
            return self.search(cur_node.right, data)
        return False

    def insert(self, cur_node, data):
        """Inserts data into tree. Alerts user if data already exists in tree.

        Calls _inspect_insertion to determine if insertion caused tree imbalance.

        :param cur_node: Root of tree.
        :param data: Data to be inserted.
        """
        if data < cur_node.data:
            if not cur_node.left:
                cur_node.left = Node(data)
                cur_node.left.parent = cur_node
                self._inspect_insertion(cur_node.left, [])
            else:
                self.insert(cur_node.left, data)

        elif data > cur_node.data:
            if not cur_node.right:
                cur_node.right = Node(data)
                cur_node.right.parent = cur_node
                self._inspect_insertion(cur_node.right, [])
            else:
                self.insert(cur_node.right, data)

        elif data == cur_node.data and cur_node.parent is not None:
            print(f'{data} already in tree. Cannot insert.')

    def _inspect_insertion(self, cur_node, nodes):
        """ Determines if insertion creates need to balance sub-tree.

        Rebalance needed if difference in height of child nodes is > 1.
        Creates list of nodes to be rotated if above condition is true.
        Calls _rebalance_node with nodes to be balanced.

        :param cur_node: The newly inserted node.
        :param nodes: Empty list.
        :return:
        """
        if not cur_node.parent:
            return

        nodes = [cur_node] + nodes

        left = self._get_height(cur_node.parent.left)
        right = self._get_height(cur_node.parent.right)

        if abs(left - right) > 1 and len(nodes) > 2:
            nodes = [cur_node.parent] + nodes
            self._rebalance_node(*nodes[:3])
            return

        new = 1 + cur_node.tallness

        if new > cur_node.parent.tallness:
            cur_node.parent.tallness = new

        self._inspect_insertion(cur_node.parent, nodes)

    def _get_height(self, cur_node):
        """ Gets height of cur_node. Returns 0 if node is None else returns node.tallness.

        :param cur_node: Node
        :return: Node.tallness
        """
        if not cur_node:
            return 0
        return cur_node.tallness

    def _rebalance_node(self, z, y, x):
        """Determines orientation of imbalanced nodes and calls indicated balancing methods.

        Calls _rotate_right or _rotate_left as determined by orientation of unbalanced nodes.

        :param z: Highest node. Rebalance occurs 'around' this node.
        :param y: Child of z
        :param x: Child of y
        """
        if y == z.left and x == y.left:
            """    z
                  /
                 y
                /
               x   """
            self._right_rotate(z)

        elif y == z.left and x == y.right:
            """   z
                 /
                y 
                 \
                  x  """
            self._left_rotate(y)
            self._right_rotate(z)

        elif y == z.right and x == y.right:
            """   z
                   \
                     y 
                      \
                        x  """
            self._left_rotate(z)

        elif y == z.right and x == y.left:
            """   z
                    \
                      y
                    /
                  x  """
            self._right_rotate(y)
            self._left_rotate(z)

        else:
            raise Exception('Tree corrupted')

    def _right_rotate(self, z):
        """Rotates around z to rebalance sub-tree.

        :param z: Root of sub-tree to be balanced.
        """
        temp = z.parent
        y = z.left
        x = y.right

        y.right = z
        z.parent = y
        z.left = x

        if x:
            x.parent = z

        y.parent = temp

        if y.parent:
            if y.parent.left == z:
                y.parent.left = y
            else:
                y.parent.right = y

        z.tallness = 1 + max(self._get_height(z.left), self._get_height(z.right))
        y.tallness = 1 + max(self._get_height(y.left), self._get_height(y.right))

    def _left_rotate(self, z):
        """Rotates around z to rebalance sub-tree.

        :param z: Root of sub-tree to be balanced.
        """
        temp = z.parent
        y = z.right
        x = y.left

        y.left = z
        z.parent = y
        z.right = x

        if x:
            x.parent = z

        y.parent = temp

        if y.parent:
            if y.parent.left == z:
                y.parent.left = y
            else:
                y.parent.right = y

        z.tallness = 1 + max(self._get_height(z.left), self._get_height(z.right))
        y.tallness = 1 + max(self._get_height(y.left), self._get_height(y.right))

    def delete(self, node):
        """ Deletes node found in _find_node.

        Removes nodes and handles deleted node's orphaned children, if any.
        Deleted nodes with two children are handled by finding the smallest relative of the deleted node's right child,
        replacing the to-be-deleted node's data with that of its smaller relative, then marking the smaller relative
        to be deleted instead. This preserves the BST imperative.
        Calls _inspect_deletion to ensure proper balancing of sub-tree after deletion.


        :param node: Node to be deleted.
        :return:
        """
        def smallest_node(curr_node):
            """ Finds smallest relative of curr_node.

            :param curr_node: A node.
            :return: Relative of curr_node with smallest value.
            """
            while curr_node.left:
                curr_node = curr_node.left
            return curr_node

        def children(curr_node):
            """ Finds number of curr_node's children.

            :param curr_node: A node
            :return: Number of curr_node's children.
            """
            num = 0
            if curr_node.left:
                num += 1
            if curr_node.right:
                num += 1
            return num

        node_parent = node.parent
        node_children = children(node)

        # Leaf nodes may simply be deleted.
        if node_children == 0:
            if node_parent:
                if node_parent.left == node:
                    node_parent.left = None
                else:
                    node_parent.right = None
            else:
                return None

        # Parent of deleted node made parent of deleted node's child.
        if node_children == 1:
            if node.left:
                child = node.left
            else:
                child = node.right
            if node_parent:
                if node_parent.left == node:
                    node_parent.left = child
                else:
                    node_parent.right = child
            else:
                child.parent = node_parent
                return child  # returned to promote child to root node
            child.parent = node_parent

        # If the node to be deleted has 2 children, the data of its smallest relative is promoted to the to-be-deleted
        # node. The smallest relative is then deleted instead.
        if node_children == 2:
            progeny = smallest_node(node.right)
            node.data = progeny.data
            self.delete(progeny)
            return

        # Adjust height and inspect the tree for balance.
        if node_parent:
            node_parent.tallness = 1 + max(self._get_height(node_parent.left), self._get_height(node_parent.right))
            self._inspect_deletion(node_parent)

    def _inspect_deletion(self, cur_node):
        """Ensures tree is balanced after deletion.

        Calls _rebalance_node if imbalance is detected.
        Calls _inspect_insertion to ensure balance up the tree.

        :param cur_node: Node. Parent of deleted node.
        :return:
        """
        if not cur_node:
            return

        left = self._get_height(cur_node.left)
        right = self._get_height(cur_node.right)

        if abs(left - right) > 1:
            y = self.taller_child(cur_node)
            x = self.taller_child(y)
            self._rebalance_node(cur_node, y, x)

        if cur_node.parent:
            self._inspect_insertion(cur_node, [])

    def taller_child(self, cur_node):
        """Finds taller of node's children.

        :param cur_node: Node. Node to be inspected.
        :return: Node. Child of curr_node with greater height.
        """
        left = self._get_height(cur_node.left)
        right = self._get_height(cur_node.right)
        if left >= right:
            return cur_node.left
        return cur_node.right


class AVLTree(object):

    """Wraps Node class. Methods call corresponding methods of Node class.

    User accessible methods for insert, delete, search, height, print_tree, and clear_tree.

    """

    def __init__(self):
        """Tree is represented by its root node, initially None."""
        self.root = None

    def __repr__(self):
        """Prints text based structure of tree.

        :return: str. Tree structure.
        """
        if not self.root:
            return 'Tree is empty. Please insert data.'
        the_tree = '\n'
        nodes = [self._get_root()]
        cur_tallness = self.root.tallness
        space = ' ' * (40 - int(len(str(self.root.data))) // 2)
        buffer = ' ' * (60 - int(len(str(self.root.data))) // 2)
        while True:
            if all(n is None for n in nodes):
                break
            cur_tallness -= 1
            this_row = ' '
            next_row = ' '
            next_nodes = []
            for cur_node in nodes:
                if not cur_node:
                    this_row += '           ' + space
                    next_nodes.extend([None, None])
                if cur_node and cur_node.data is not Node:
                    this_row += f'{buffer}{str(cur_node.data)}{buffer}'
                if cur_node and cur_node.left:
                    next_nodes.append(cur_node.left)
                    next_row += space + '/' + space
                else:
                    next_nodes.append(None)
                    next_row += '       ' + space
                if cur_node and cur_node.right:
                    next_nodes.append(cur_node.right)
                    next_row += '\\' + space
                else:
                    next_nodes.append(None)
                    next_row += '       ' + space
            the_tree += (cur_tallness * '   ' + this_row + '\n' + cur_tallness * '   ' + next_row + '\n')
            space = ' ' * int(len(space) // 2)
            buffer = ' ' * int(len(buffer) // 2)
            nodes = next_nodes
        return the_tree

    def _get_root(self, data=None):
        """Returns root node.

        Creates root node if called from Tree.insert and tree is empty.

        :param data: If provided and tree is empty, tree root is established with data.
        :return: Root node of tree.
        """
        if not self.root:
            if data is not None:
                self.root = Node(data)
            else:
                print('Tree is empty.')
                return

        else:
            while self.root.parent:
                self.root = self.root.parent

        return self.root

    def print_tree(self, order=None):
        """User interface for printing tree.

        Calls _print_tree with root from _get_root.
        Prints order as entered by user for correcting typos and such.
        Ensures print order requested is valid.

        :param order: 'in-order', 'pre-order', or 'post-order'
        :return: _print_tree
        """
        print(order)
        if not order or not any([order == 'pre-order', order == 'post-order', order == 'in-order']):
            print('Please specify a valid print order.')
            print('(eg. order=\'in-order\', \'pre-order\', or \'post-order\')')
            return
        return self._print_tree(self._get_root(), order)

    def _print_tree(self, root, order=None):
        """Calls print_tree method of Node class.

        :param root: Root node
        :param order: 'in-order', 'pre-order', or 'post-order'
        :return: Node.print_tree
        """
        return root.print_tree(root, order)

    def height(self, print_result=False):
        """User interface for finding height of tree.

        Calls _height with root from _get_root.
        Option to print height to stdout.

        :param print_result: Prints height to stdout if True.
        :return: _height
        """
        height = self._height(self._get_root())
        if print_result:
            print(height)
        return height

    def _height(self, root):
        """Calls height method of Node class.

        :param root: Root node.
        :return: Node.height
        """
        return root.height(root, 0)

    def search(self, data, print_result=False):
        """User interface for search method.

        Calls _search with root from _get_root.
        Option to print results to stdout.

        :param print_result: Set to True to print search results.
        :param data: Data to be found in tree.
        :return: _search
        """
        result = self._search(self._get_root(), data)

        if print_result:
            if result:
                print(f'{data} found.')
            else:
                print(f'{data} not found.')

        return result

    def _search(self, root, data):
        """Calls search method of Node class.

        :param root: Root node.
        :param data: Data to search for in tree.
        :return: Node.search
        """
        return root.search(root, data)

    def insert(self, data):
        """User interface for inserting data into tree.

        Calls _insert with root provided by _get_root.
        Data field provided to _get_root ensures tree root is created on first insertion.

        :param data: Data to be inserted into tree.
        :return: _insert
        """
        return self._insert(self._get_root(data=data), data)

    def _insert(self, root, data):
        """Calls insert method of Node class.

        :param root: Root node.
        :param data: Data to be inserted into tree.
        :return: Node.insert
        """
        return root.insert(root, data)


    def _find_node(self, cur_node, data):
        """Finds and returns node with given data, else, returns None.

        :param cur_node: Root node from _get_root.
        :param data: Data contained within node to be found.
        :return: Node containing data if such a node exists. Else, None.
        """
        if cur_node and data == cur_node.data:
            return cur_node
        elif cur_node and data < cur_node.data:
            return self._find_node(cur_node.left, data)
        elif cur_node and data > cur_node.data:
            return self._find_node(cur_node.right, data)
        return None

    def delete(self, data):
        """Passes root and data to be deleted to _delete.

        Calls _delete with root from _get_root.

        :param data: Data to delete from tree.
        :return: _delete if data in tree, else, None.
        """
        node = self._find_node(self._get_root(), data)

        if not node:
            print(f'{data} not in tree, Cannot delete.')
            return

        return self._delete(node)

    def _delete(self, node):
        """Calls delete method of Node class.

        If root is to be deleted and has no children, root is set to None.
        If Node.delete returns a node it is the new root node.

        :param node: Node to be deleted
        """
        if node == self.root and (not node.left and not node.right):
            self.root = None
            return self.root

        result = node.delete(node)

        if result:
            self.root = result

        return self.root

    def clear_tree(self):
        """Clears tree of all data.

        :return: Tree root
        """
        self.root = None
        return self.root

Thank you for taking a look!

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.