Thursday, September 15, 2011

AVL - Binary balanced tree in C#

This is my implementation of a binary search tree with both recursive and iterative insert:
 
using System;
using System.Collections.Generic;
using System.Diagnostics;

namespace AVLTree
{
	public class AVLTree<ElementType>
		where ElementType : IComparable<ElementType>
	{
		public Node Root { get; set; }
		public Action<Node, ElementType> PerformActionOnInsertingNodesAllreadyInTheTree { get; set; }

		public void Insert(ElementType value)
		{
			Root = Insert(value, Root);
		}

		public IEnumerable<ElementType> GetContent()
		{
			foreach (var item in GetContent(Root))
				yield return item.Value;
		}

		public int GetHeight()
		{
			int max = int.MinValue;
			foreach (var item in GetContent(Root))
				max = Math.Max(max, item.Height);
			return max;
		}

		private IEnumerable<Node> GetContent(Node node)
		{
			if (null == node)
				yield break;

			foreach (var item in GetContent(node.Left))
				yield return item;
			yield return node;
			foreach (var item in GetContent(node.Right))
				yield return item;
		}

		private Node Insert(ElementType value, Node subtreeNode)
		{
			if (null == subtreeNode)
			{
				subtreeNode = new Node(value);
				return subtreeNode;
			}
			if (0 > value.CompareTo(subtreeNode.Value))
			{
				subtreeNode.Left = Insert(value, subtreeNode.Left);
				if (1 < subtreeNode.Left.Height - Node.HeightOf(subtreeNode.Right))
					if (0 > value.CompareTo(subtreeNode.Left.Value))
						subtreeNode = subtreeNode.RotateWithLeftChild();
					else
						subtreeNode = subtreeNode.DoubleRotateWithLeftChild();
			}
			else if (0 < value.CompareTo(subtreeNode.Value))
			{
				subtreeNode.Right = Insert(value, subtreeNode.Right);
				if (1 < subtreeNode.Right.Height - Node.HeightOf(subtreeNode.Left))
					if (0 < value.CompareTo(subtreeNode.Right.Value))
						subtreeNode = subtreeNode.RotateWithRightChild();
					else
						subtreeNode = subtreeNode.DoubleRotateWithRightChild();
			}
			else
			{
				if (null != PerformActionOnInsertingNodesAllreadyInTheTree)
					PerformActionOnInsertingNodesAllreadyInTheTree(subtreeNode, value);
			}

			subtreeNode.Height = Math.Max(Node.HeightOf(subtreeNode.Left), Node.HeightOf(subtreeNode.Right)) + 1;
			return subtreeNode;
		}

		public void InsertPlain(ElementType value)
		{
			Node current = Root;
			Stack<TreeNavigationStep> pathToCurrentNode = new Stack<TreeNavigationStep>();

			while (null != current)
			{
				if (0 > value.CompareTo(current.Value))
				{
					pathToCurrentNode.Push(new TreeNavigationStep(current, true, true));
					current = current.Left;
				}
				else if (0 < value.CompareTo(current.Value))
				{
					pathToCurrentNode.Push(new TreeNavigationStep(current, false, true));
					current = current.Right;
				}
				else
				{
					if (null != PerformActionOnInsertingNodesAllreadyInTheTree)
					{
						PerformActionOnInsertingNodesAllreadyInTheTree(current, value);
					}
					return;
				}
			}
			Node newNode = new Node(value);
			if (0 == pathToCurrentNode.Count)
			{
				Root = newNode;
				return;
			}

			TreeNavigationStep navigationStep = pathToCurrentNode.Peek();
			current = navigationStep.From;

			if (navigationStep.ToLeft)
			{
				current.Left = newNode;
			}
			else
			{
				current.Right = newNode;
			}

			while(pathToCurrentNode.Count >0)
			{
				var step = pathToCurrentNode.Pop();
				bool performedBalance = false;
				Node parent = step.From;
				Node nodeResultedByBalance = null;
				if (step.FreeToBalance)
				{
					if (step.ToLeft)
					{
						if (1 < parent.Left.Height - Node.HeightOf(parent.Right))
							if (0 > value.CompareTo(parent.Left.Value))
							{
								performedBalance = true;
								nodeResultedByBalance = parent.RotateWithLeftChild();
								if(0<pathToCurrentNode.Count)
								{
									var gradpa = pathToCurrentNode.Peek();
									if(gradpa.ToLeft)
										gradpa.From.Left = nodeResultedByBalance;
									else
										gradpa.From.Right = nodeResultedByBalance;
								}
								pathToCurrentNode.Push(new TreeNavigationStep(nodeResultedByBalance, false, false));
							}
							else
							{
								performedBalance = true;
								nodeResultedByBalance = parent.DoubleRotateWithLeftChild();
								if(0<pathToCurrentNode.Count)
								{
									var gradpa = pathToCurrentNode.Peek();
									if(gradpa.ToLeft)
										gradpa.From.Left = nodeResultedByBalance;
									else
										gradpa.From.Right = nodeResultedByBalance;
								}
								pathToCurrentNode.Push(new TreeNavigationStep(nodeResultedByBalance.Left, false, false));
								pathToCurrentNode.Push(new TreeNavigationStep(nodeResultedByBalance, false, false));
							}
					}
					else
					{
						if (1 < parent.Right.Height - Node.HeightOf(parent.Left))
							if (0 < value.CompareTo(parent.Right.Value))
							{
								performedBalance = true;
								nodeResultedByBalance = parent.RotateWithRightChild();
								if(0<pathToCurrentNode.Count)
								{
									var gradpa = pathToCurrentNode.Peek();
									if(gradpa.ToLeft)
										gradpa.From.Left = nodeResultedByBalance;
									else
										gradpa.From.Right = nodeResultedByBalance;
								}
								pathToCurrentNode.Push(new TreeNavigationStep(nodeResultedByBalance, false, false));
							}
							else
							{
								performedBalance = true;
								nodeResultedByBalance = parent.DoubleRotateWithRightChild();
								if(0<pathToCurrentNode.Count)
								{
									var gradpa = pathToCurrentNode.Peek();
									if(gradpa.ToLeft)
										gradpa.From.Left = nodeResultedByBalance;
									else
										gradpa.From.Right = nodeResultedByBalance;
								}
								pathToCurrentNode.Push(new TreeNavigationStep(nodeResultedByBalance.Right, false, false));
								pathToCurrentNode.Push(new TreeNavigationStep(nodeResultedByBalance, false, false));
							}
					}
					if (performedBalance)
						Root = nodeResultedByBalance;
					else
						Root = parent;
				}
				
				parent.Height = Math.Max(Node.HeightOf(parent.Left), Node.HeightOf(parent.Right)) + 1;
			}
		}

		public class TreeNavigationStep
		{
			public TreeNavigationStep(Node from, bool toLeftChild, bool freeToBalance)
			{
				From = from;
				ToLeft = toLeftChild;
				FreeToBalance = freeToBalance;
			}
			public Node From { get; set; }
			public bool ToLeft { get; set; }
			public bool FreeToBalance { get; set; }
		}

		[DebuggerDisplay("{Value}")]
		public class Node
		{
			public Node(ElementType value)
			{
				Value = value;
			}

			public Node Left { get; set; }
			public Node Right { get; set; }
			public ElementType Value { get; set; }
			public int Height { get; set; }

			public Node RotateWithLeftChild()
			{
				Debug.Assert(null != this.Left);

				Node left = this.Left;
				this.Left = left.Right;
				left.Right = this;

				this.Height = Math.Max(HeightOf(this.Left), HeightOf(this.Right)) + 1;
				left.Height = Math.Max(HeightOf(left.Left), this.Height) + 1;
				return left;
			}

			public Node RotateWithRightChild()
			{
				Debug.Assert(null != this.Right);

				Node right = this.Right;
				this.Right = right.Left;
				right.Left = this;

				this.Height = Math.Max(HeightOf(this.Left), HeightOf(this.Right)) + 1;
				right.Height = Math.Max(HeightOf(right.Right), this.Height) + 1;
				return right;
			}

			public Node DoubleRotateWithLeftChild()
			{
				Debug.Assert(null != this.Left && null != this.Left.Right);

				this.Left = this.Left.RotateWithRightChild();
				return this.RotateWithLeftChild();
			}

			public Node DoubleRotateWithRightChild()
			{
				Debug.Assert(null != this.Right && null != this.Right.Left);

				this.Right = this.Right.RotateWithLeftChild();
				return this.RotateWithRightChild();
			}

			public static int HeightOf(Node node)
			{
				if (null == node)
					return -1;
				return node.Height;
			}
		}
	}
}