I decided to take it upon myself to create a Binary Search Tree balancing function. Although there may be a methodology like the delete function where you can just look at a couple nodes at a time, the ripple effects could potentially be unpredictable. And although you may be able to balance the tree a bit at a time, it would probably be slower overall than to just do the whole tree at once.

 

So, let’s think about a slow balance feature (probably far from optimal).

1. Count nodes to left of tree and right of the tree
2. If the left tree is larger, take the greatest value in the left tree and replace the root with that value, then place the root value at the appropriate spot in the right tree.
3. If the right tree is larger, take the smallest value in the right tree and replace the root with that value, then place the root value at the appropriate spot in the left tree.
4. Repeat steps above until the tree is perfectly balanced or off by 1.
5. Congrats, your tree is balanced at the root! But what about elsewhere? Your tree of n nodes could just be a V shape and still incredibly inefficient!
6. So now that the root is balanced, recursively call the function that does steps 1-3 on root.left and root.right!

This should work. And this you should be able to halt at any time between iterations and have made progress. But to have it completely balance a tree sounds like it would be pretty inefficient since it rebuilds the tree one node at a time.

What I did instead was create a full balance function that requires you to be able to rebuild the whole tree in one sitting. It’s totally fine given the size trees we expect to work with. The algorithm works as follows:

1. Create a sorted list of all values in the tree
2. Call a function that takes the list as an argument.
3. Build a new tree based on this value with the middle element at the root and recursively call the function on [:mid] for the left node and [mid+1:] for the right node.

It works and I think it’s pretty cool! Here’s the code (assuming a functional insert function and a BinaryTree class).

note: since doing the puzzle assignment, I’ve taken a liking to using nested functions for recursive functions that require the maintenance of a data structure. It’s so much harder to make a mistake doing it.

def list_of_values(t):
    """
    Return a sorted list of values at tree rooted at node t

    @type t: BinaryTree
    @rtype: list

    >>> b = BinaryTree(1)
    >>> b = insert(b, 12)
    >>> b = insert(b, 3)
    >>> list_of_values(b)
    [1, 3, 12]
    """
    nodes = set()

    def list_nodes(tree):
        nodes.add(tree.data)
        if tree.left:
            nodes.add(list_nodes(tree.left))
        if tree.right:
            nodes.add(list_nodes(tree.right))

    list_nodes(t)
    ordered_nodes = list(nodes - {None})
    ordered_nodes.sort()
    return ordered_nodes


def balance_tree(t):
    """
    Balance and return the binary search tree rooted at node t

    @type t: BinaryTree
    @rtype: BinaryTree

    >>> b = BinaryTree(7)
    >>> b = insert(b, 6)
    >>> b = insert(b, 5)
    >>> b = insert(b, 4)
    >>> b = insert(b, 3)
    >>> b = insert(b, 1)
    >>> b = insert(b, 0)
    >>> print(balance_tree(b))
            7
        6
            5
    4
            3
        1
            0
    <BLANKLINE>
    >>> b = insert(b, 14)
    >>> b = insert(b, 28)
    >>> b = insert(b, 2)
    >>> print(balance_tree(b))
            28
        14
            7
                6
    5
            4
                3
        2
            1
                0
    <BLANKLINE>
    """
    ordered_tree_data = list_of_values(t)

    def new_tree(l):
        if not l:
            return None
        mid = len(l)//2
        return BinaryTree(l[mid], new_tree(l[:mid]), new_tree(l[mid+1:]))

    return new_tree(ordered_tree_data)