Add index operations for AVL tree

Author: justcoding121

src/Advanced.Algorithms/DataStructures/Tree/AvlTree.cs

@@ -33,7 +33,7 @@ public AVLTree(bool enableNodeLookUp = false)
         /// <param name="sortedKeys">The sorted keys.</param>
         /// <param name="enableNodeLookUp">Enabling lookup will fasten deletion/insertion/exists operations
         ///  at the cost of additional space.</param>
-        public AVLTree(IEnumerable<T> sortedKeys, bool enableNodeLookUp = false) 
+        public AVLTree(IEnumerable<T> sortedKeys, bool enableNodeLookUp = false)
             : this(enableNodeLookUp)
         {
             ValidateCollection(sortedKeys);
@@ -43,7 +43,7 @@ public AVLTree(IEnumerable<T> sortedKeys, bool enableNodeLookUp = false)
             assignCount(Root);
         }
 
-     
+
         /// <summary>
         /// Time complexity: O(log(n))
         /// </summary>
@@ -139,6 +139,25 @@ private void insert(AVLTreeNode<T> node, T value)
             node.UpdateCounts();
         }
 
+        /// <summary>
+        /// </summary>
+        public int IndexOf(T item)
+        {
+            return Root.Position(item);
+        }
+
+        /// <summary>
+        ///  Time complexity: O(log(n))
+        /// </summary>
+        public T ElementAt(int index)
+        {
+            if (index < 0 || index >= Count)
+            {
+                throw new ArgumentNullException("index");
+            }
+
+            return Root.KthSmallest(index).Value;
+        }
 
         /// <summary>
         /// Time complexity: O(log(n)).
@@ -159,8 +178,37 @@ public void Delete(T value)
 
         }
 
-        
-        private void delete(AVLTreeNode<T> node, T value)
+        /// <summary>
+        ///  Time complexity: O(log(n))
+        /// </summary>
+        public T RemoveAt(int index)
+        {
+            if (index < 0 || index >= Count)
+            {
+                throw new ArgumentException("index");
+            }
+
+            var nodeToDelete = Root.KthSmallest(index) as AVLTreeNode<T>;
+            var nodeToBalance = delete(nodeToDelete, nodeToDelete.Value);
+
+            while (nodeToBalance != null)
+            {
+                nodeToBalance.UpdateCounts();
+                updateHeight(nodeToBalance);
+                balance(nodeToBalance);
+
+                nodeToBalance = nodeToBalance.Parent;
+            }
+
+            if (nodeLookUp != null)
+            {
+                nodeLookUp.Remove(nodeToDelete.Value);
+            }
+
+            return nodeToDelete.Value;
+        }
+
+        private AVLTreeNode<T> delete(AVLTreeNode<T> node, T value)
         {
             var baseCase = false;
 
@@ -289,14 +337,15 @@ private void delete(AVLTreeNode<T> node, T value)
                 node.Parent.UpdateCounts();
                 updateHeight(node.Parent);
                 balance(node.Parent);
+                return node.Parent;
             }
             else
             {
                 node.UpdateCounts();
                 updateHeight(node);
                 balance(node);
+                return node;
             }
-
         }
 
         /// <summary>
@@ -307,7 +356,6 @@ public T FindMax()
             return findMax(Root).Value;
         }
 
-
         private AVLTreeNode<T> findMax(AVLTreeNode<T> node)
         {
             while (true)
@@ -565,7 +613,7 @@ private void updateHeight(AVLTreeNode<T> node)
 
         private void recomputeHeight(AVLTreeNode<T> node)
         {
-            if(node == null)
+            if (node == null)
             {
                 return;
             }
@@ -629,6 +677,19 @@ internal void Swap(T value1, T value2)
             }
         }
 
+        /// <summary>
+        /// Descending enumerable.
+        /// </summary>
+        public IEnumerable<T> AsEnumerableDesc()
+        {
+            return GetEnumeratorDesc().AsEnumerable();
+        }
+
+        public IEnumerator<T> GetEnumeratorDesc()
+        {
+            return new BSTEnumerator<T>(Root, false);
+        }
+
         //Implementation for the GetEnumerator method.
         IEnumerator IEnumerable.GetEnumerator()
         {

src/Advanced.Algorithms/DataStructures/Tree/RedBlackTree.cs

@@ -381,7 +381,7 @@ public T RemoveAt(int index)
         {
             if (index < 0 || index >= Count)
             {
-                throw new ArgumentNullException("index");
+                throw new ArgumentException("index");
             }
 
             var node = Root.KthSmallest(index) as RedBlackTreeNode<T>;

tests/Advanced.Algorithms.Tests/DataStructures/Tree/AVLTree_Tests.cs

@@ -99,17 +99,18 @@ public void AVLTree_Accuracy_Test()
             var nodeCount = 1000;
 
             var rnd = new Random();
-            var sortedNumbers = Enumerable.Range(1, nodeCount)
+            var sorted = Enumerable.Range(1, nodeCount).ToList();
+            var randomNumbers = sorted
                                 .OrderBy(x => rnd.Next())
                                 .ToList();
 
             var tree = new AVLTree<int>(true);
 
             for (int i = 0; i < nodeCount; i++)
             {
-                tree.Insert(sortedNumbers[i]);
+                tree.Insert(randomNumbers[i]);
 
-                Assert.IsTrue(tree.HasItem(sortedNumbers[i]));
+                Assert.IsTrue(tree.HasItem(randomNumbers[i]));
                 Assert.IsTrue(tree.Root.IsBinarySearchTree(int.MinValue, int.MaxValue));
                 tree.Root.VerifyCount();
 
@@ -122,19 +123,35 @@ public void AVLTree_Accuracy_Test()
                 Assert.IsTrue(tree.Count == i + 1);
             }
 
-            sortedNumbers = Enumerable.Range(1, nodeCount)
+            for (int i = 0; i < sorted.Count; i++)
+            {
+                Assert.AreEqual(sorted[i], tree.ElementAt(i));
+                Assert.AreEqual(i, tree.IndexOf(sorted[i]));
+            }
+
+            randomNumbers = Enumerable.Range(1, nodeCount)
                                 .OrderBy(x => rnd.Next())
                                 .ToList();
 
             //IEnumerable test using linq
             Assert.AreEqual(tree.Count, tree.Count());
+            Assert.AreEqual(tree.Count, tree.AsEnumerableDesc().Count());
 
             for (int i = 0; i < nodeCount; i++)
             {
-                tree.Delete(sortedNumbers[i]);
+                if (rnd.NextDouble() >= 0.5)
+                {
+                    tree.Delete(randomNumbers[i]);
+                }
+                else
+                {
+                    var index = tree.IndexOf(randomNumbers[i]);
+                    Assert.AreEqual(tree.ElementAt(index), randomNumbers[i]);
+                    tree.RemoveAt(index);
+                }
 
-                tree.Root.VerifyCount();
                 Assert.IsTrue(tree.Root.IsBinarySearchTree(int.MinValue, int.MaxValue));
+                tree.Root.VerifyCount();
 
                 var actualHeight = tree.GetHeight();
 

tests/Advanced.Algorithms.Tests/DataStructures/Tree/RedBlackTree_Tests.cs

@@ -87,7 +87,6 @@ public void RedBlackTree_Accuracy_Test()
                 Assert.AreEqual(i, tree.IndexOf(sorted[i]));
             }
 
-
             //shuffle again before deletion tests
             randomNumbers = Enumerable.Range(1, nodeCount)
                                    .OrderBy(x => rnd.Next())
@@ -106,6 +105,7 @@ public void RedBlackTree_Accuracy_Test()
                 else
                 {
                     var index = tree.IndexOf(randomNumbers[i]);
+                    Assert.AreEqual(tree.ElementAt(index), randomNumbers[i]);
                     tree.RemoveAt(index);
                 }