enable node lookUp for R-B tree

Author: justcoding121

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

@@ -17,27 +17,40 @@ public class RedBlackTree<T> : BSTBase<T>, IEnumerable<T> where T : IComparable
 
         public int Count => Root == null ? 0 : Root.Count;
 
+        /// <param name="enableNodeLookUp">Enabling lookup will fasten deletion/insertion/exists operations
+        /// at the cost of additional space.</param>
+        public RedBlackTree(bool enableNodeLookUp = false)
+        {
+            if (enableNodeLookUp)
+            {
+                nodeLookUp = new Dictionary<T, BSTNodeBase<T>>();
+            }
+        }
+
         /// <summary>
         /// Initialize the BST with given sorted keys optionally.
         /// Time complexity: O(n).
         /// </summary>
         /// <param name="sortedCollection">The sorted initial collection.</param>
-        public RedBlackTree(IEnumerable<T> sortedCollection = null)
+        /// <param name="enableNodeLookUp">Enabling lookup will fasten deletion/insertion/exists operations
+        /// at the cost of additional space.</param>
+        public RedBlackTree(IEnumerable<T> sortedCollection, bool enableNodeLookUp = false)
         {
-            if (sortedCollection != null)
+            ValidateSortedCollection(sortedCollection);
+            var nodes = sortedCollection.Select(x => new RedBlackTreeNode<T>(null, x)).ToArray();
+            Root = (RedBlackTreeNode<T>)ToBST(nodes);
+            assignColors(Root);
+            assignCount(Root);
+
+            if (enableNodeLookUp)
             {
-                ValidateSortedCollection(sortedCollection);
-                var nodes = sortedCollection.Select(x => new RedBlackTreeNode<T>(null, x)).ToArray();
-                Root = (RedBlackTreeNode<T>)ToBST(nodes);
-                assignColors(Root);
-                assignCount(Root);
+                nodeLookUp = nodes.ToDictionary(x => x.Value, x => x as BSTNodeBase<T>);
             }
         }
 
         ///Special (internal only) constructor for Bentley-Ottmann sweep line algorithm for fast line swap.
         /// <param name="equalityComparer">Provide custom IEquality comparer for node lookup dictionary.</param>
         internal RedBlackTree(IEqualityComparer<T> equalityComparer)
-            : this()
         {
             nodeLookUp = new Dictionary<T, BSTNodeBase<T>>(equalityComparer);
         }
@@ -131,9 +144,13 @@ internal bool Exists(T value)
         {
             if (nodeLookUp != null)
             {
-                //since node look up is only used by Bentley-Ottmann algorithm internally
-                //and it does'nt need the position we can return defualt(int).
-                return (nodeLookUp[value] as RedBlackTreeNode<T>, default(int));
+                if (nodeLookUp.ContainsKey(value))
+                {
+                    var node = (nodeLookUp[value] as RedBlackTreeNode<T>);
+                    return (node, Root.Position(value));
+                }
+
+                return (null, -1);
             }
 
             var result = Root.Find(value);
@@ -391,11 +408,13 @@ public T RemoveAt(int index)
 
             var node = Root.KthSmallest(index) as RedBlackTreeNode<T>;
 
+            var deletedValue = node.Value;
+
             delete(node);
 
             if (nodeLookUp != null)
             {
-                nodeLookUp.Remove(node.Value);
+                nodeLookUp.Remove(deletedValue);
             }
 
             return node.Value;
@@ -439,13 +458,14 @@ private void delete(RedBlackTreeNode<T> node)
                 {
                     var maxLeftNode = findMax(node.Left);
 
-                    node.Value = maxLeftNode.Value;
-
                     if (nodeLookUp != null)
                     {
-                        nodeLookUp[node.Value] = node;
+                        nodeLookUp[node.Value] = maxLeftNode;
+                        nodeLookUp[maxLeftNode.Value] = node;
                     }
 
+                    node.Value = maxLeftNode.Value;
+
                     //delete left max node
                     delete(maxLeftNode);
                     return;

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

@@ -125,6 +125,78 @@ public void RedBlackTree_Accuracy_Test()
             Assert.IsTrue(tree.Count == 0);
         }
 
+        [TestMethod]
+        public void RedBlack_Accuracy_Test_With_Node_LookUp()
+        {
+            var nodeCount = 1000;
+
+            var rnd = new Random();
+            var sorted = Enumerable.Range(1, nodeCount).ToList();
+            var randomNumbers = sorted
+                                .OrderBy(x => rnd.Next())
+                                .ToList();
+
+            var tree = new RedBlackTree<int>(true);
+
+            for (int i = 0; i < nodeCount; i++)
+            {
+                var index = tree.Insert(randomNumbers[i]);
+                Assert.AreEqual(index, tree.IndexOf(randomNumbers[i]));
+                Assert.IsTrue(tree.HasItem(randomNumbers[i]));
+                Assert.IsTrue(tree.Root.IsBinarySearchTree(int.MinValue, int.MaxValue));
+                tree.Root.VerifyCount();
+                var actualHeight = tree.Root.GetHeight();
+
+                //http://doctrina.org/maximum-height-of-red-black-tree.html
+                var maxHeight = 2 * Math.Log(nodeCount + 1, 2);
+
+                Assert.IsTrue(actualHeight < maxHeight);
+                Assert.IsTrue(tree.Count == i + 1);
+            }
+
+            for (int i = 0; i < sorted.Count; i++)
+            {
+                Assert.AreEqual(sorted[i], tree.ElementAt(i));
+                Assert.AreEqual(i, tree.IndexOf(sorted[i]));
+            }
+
+            //shuffle again before deletion tests
+            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++)
+            {
+                if (rnd.NextDouble() >= 0.5)
+                {
+                    var index = tree.IndexOf(randomNumbers[i]);
+                    Assert.AreEqual(index, tree.Delete(randomNumbers[i]));
+                }
+                else
+                {
+                    var index = tree.IndexOf(randomNumbers[i]);
+                    Assert.AreEqual(tree.ElementAt(index), randomNumbers[i]);
+                    tree.RemoveAt(index);
+                }
+
+                Assert.IsTrue(tree.Root.IsBinarySearchTree(int.MinValue, int.MaxValue));
+                tree.Root.VerifyCount();
+                var actualHeight = tree.Root.GetHeight();
+
+                //http://doctrina.org/maximum-height-of-red-black-tree.html
+                var maxHeight = 2 * Math.Log(nodeCount + 1, 2);
+
+                Assert.IsTrue(actualHeight < maxHeight);
+                Assert.IsTrue(tree.Count == nodeCount - 1 - i);
+            }
+
+            Assert.IsTrue(tree.Count == 0);
+        }
+
         [TestMethod]
         public void RedBlackTree_BulkInit_Test()
         {
@@ -162,6 +234,44 @@ public void RedBlackTree_BulkInit_Test()
             Assert.IsTrue(tree.Count == 0);
         }
 
+
+        [TestMethod]
+        public void RedBlackTree_BulkInit_Test_With_Node_LookUp()
+        {
+            var nodeCount = 1000;
+
+            var rnd = new Random();
+            var sortedNumbers = Enumerable.Range(1, nodeCount).ToList();
+
+            var tree = new RedBlackTree<int>(sortedNumbers, true);
+
+            Assert.IsTrue(tree.Root.IsBinarySearchTree(int.MinValue, int.MaxValue));
+
+            //IEnumerable test using linq
+            Assert.AreEqual(tree.Count, tree.Count());
+            Assert.AreEqual(tree.Count, tree.AsEnumerableDesc().Count());
+
+            tree.Root.VerifyCount();
+
+            for (int i = 0; i < nodeCount; i++)
+            {
+                tree.Delete(sortedNumbers[i]);
+
+                tree.Root.VerifyCount();
+                Assert.IsTrue(tree.Root.IsBinarySearchTree(int.MinValue, int.MaxValue));
+
+                var actualHeight = tree.Root.GetHeight();
+
+                //http://doctrina.org/maximum-height-of-red-black-tree.html
+                var maxHeight = 2 * Math.Log(nodeCount + 1, 2);
+
+                Assert.IsTrue(actualHeight < maxHeight);
+                Assert.IsTrue(tree.Count == nodeCount - 1 - i);
+            }
+
+            Assert.IsTrue(tree.Count == 0);
+        }
+
         [TestMethod]
         public void RedBlackTree_StressTest()
         {