Balanced tree with node rotation on insert

Adding values to a binary tree does not guarentee a balanced tree. For example, a root of M into which A, Z, D, F, E are sequentially added ends up as

             M
           / \
          A   Z
           \
            D
             \
              F
             /
            E
        

The trick is to check on the insert of a new node whether the tree from the grandparent down is 'balanced' or not. Starting at the granparent, walk the tree subtracting 1 when the left is populated and adding one when the right is populated. If the result is -1,0, or +1 the tree is balanced. If the result is 2 or -2 then we need to perform a rotation to bring the 'mid-point' node back into position. If we start with a balanced tree, and do this on every insert, there is no way for the tree to become unbalanced.

-2 -1 0 1 2
Left heavy, rotate right so that
       C                 B         
      /                 / \         
     B      becomes    A   C        
    /                               
   A                                
   
Balanced. No change needed
   C  
 /   
A    
Balanced. No change needed
    C
  /
A
  \
    B
Balanced. No change needed
     B            
    / \            
     A   C  no change
Balanced. No change needed
A
 \
  C
 /
B
Right heavy, rotate left so that
 C                 B       
  \               / \       
 B    becomes  A   C    
   \                      
   A                    
class Program
{
    static void Main(string[] args)
    {
        Tree tree = new Tree();
        tree.Add(new Node("M""abc"));
        tree.Add(new Node("Z""zzz"));
        tree.Add(new Node("D""ddd"));
        tree.Add(new Node("A""aaa"));
        tree.Add(new Node("E""eee"));
        tree.Add(new Node("F""fff"));
        tree.Add(new Node("G""ggg"));
        tree.Add(new Node("EF""eeff"));
 
        Node f = tree.Find("F");
        Console.WriteLine(f.Payload);
        Node b = tree.Find("B");
        Console.WriteLine(b == null);
 
        Console.WriteLine("F {0}", tree.Find("F").CalculateBalance());
        Console.WriteLine("D  {0}", tree.Find("D").CalculateBalance());
        Console.WriteLine("G {0}", tree.Find("G").CalculateBalance());
 
        foreach (Node n in tree.Root.Walk())
            Console.WriteLine(n.Key);
    }
}
 
class Tree
{
    public Node Root;
 
    public void Add(Node newNode)
    {
        if (this.Root == null)
            this.Root = newNode;
        else
            this.Add(this.Root, newNode);
    }
 
    public Node Find(string key)
    {
        return this.Find(this.Root, key);
    }
 
    private Node Find(Node thisNode, string key)
    {
        int compare = string.Compare(thisNode.Key, thisNode.Key);
        if (compare == 0) return thisNode;
        if (compare < 0)
        {
            if (thisNode.Left == nullreturn null;
            return this.Find(thisNode.Left, key);
        }
        if (thisNode.Right == nullreturn null;
        return this.Find(thisNode.Right, key);
    }
 
    public void Add(Node thisNode, Node newNode)
    {
        newNode.Parent = thisNode;
        int compare = string.Compare(newNode.Key, thisNode.Key);
        // if (compare == 0) return this;
        if (compare < 0)
        {
            if (thisNode.Left == null)
                thisNode.Left = newNode;
            else
                this.Add(thisNode.Left, newNode);
        }
        else
        {
            if (thisNode.Right == null)
                thisNode.Right = newNode;
            else
                this.Add(thisNode.Right, newNode);
        }
 
        // Did we unbalance the tree?
        if (newNode.Parent == thisNode)
        {
            Node thisParent = thisNode.Parent;
            if (thisParent == nullreturn;
            bool replaceRoot = thisParent.Parent == null;
            int balance = thisNode.Parent.CalculateBalance();
            Node a, b, c;
            if (balance < -1)
            {
                // Left heavy, rotate right so that
                //    C                 B
                //   B      becomes    A C
                //  A
                c = thisNode.Parent;
                b = thisNode;
                a = newNode;
                if (c.Parent != null)
                {
                    if (c.Parent.Left == c)
                        c.Parent.Left = b;
                    else
                        c.Parent.Right = b;
                }
                b.Parent = c.Parent;
                c.Parent = b;
                a.Parent = a;
                c.Left = null;
                b.Right = c;
 
                if (replaceRoot) this.Root = b;
            }
            else if (balance > 1)
            {
                // Right heavy, rotate left so that
                //  A                   B
                //   B      becomes    A C
                //    C
                a = thisNode.Parent;
                b = thisNode;
                c = newNode;
                if (a.Parent != null)
                {
                    if (a.Parent.Left == a)
                        a.Parent.Left = b;
                    else
                        a.Parent.Right = b;
                }
                a.Parent = b;
                c.Parent = b;
                b.Left = a;
                b.Right = c;
                a.Right = null;
 
                if (replaceRoot) this.Root = b;
            }
        }
    }
}
 
public class Node
{
    public object Payload { getset; }
    public string Key { getinternal set; }
    public Node Left { getinternal set; }
    public Node Right { getinternal set; }
    public Node Parent { getinternal set; }
 
    public Node(string key, object payload)
    {
        this.Key = key;
        this.Payload = payload;
    }
 
    public IEnumerable<Node> Walk()
    {
        if (this.Left != nullforeach (Node n in this.Left.Walk()) yield return n;
        yield return this;
        if (this.Right != nullforeach (Node n in this.Right.Walk()) yield return n;
    }
 
    public int CalculateBalance()
    {
        int balance = 0;
        if (this.Left != null)
        {
            balance--;
            balance += this.Left.CalculateBalance();
        }
        if (this.Right != null)
        {
            balance++;
            balance += this.Right.CalculateBalance();
        }
        return balance;
    }
}