using System; using System.Collections; using System.Collections.Generic; using System.Linq; namespace Advanced.Algorithms.DataStructures { /// /// A red black tree implementation. /// public class RedBlackTree : BSTBase, IEnumerable where T : IComparable { private readonly Dictionary> nodeLookUp; internal RedBlackTreeNode Root { get; set; } public int Count => Root == null ? 0 : Root.Count; /// Enabling lookup will fasten deletion/insertion/exists operations /// at the cost of additional space. /// Provide custom IEquality comparer for node lookup dictionary when enabled. public RedBlackTree(bool enableNodeLookUp = false, IEqualityComparer equalityComparer = null) { if (enableNodeLookUp) { nodeLookUp = new Dictionary>(equalityComparer); } } /// /// Initialize the BST with given sorted keys. /// Time complexity: O(n). /// /// The sorted keys. /// Enabling lookup will fasten deletion/insertion/exists operations /// at the cost of additional space. /// Provide custom IEquality comparer for node lookup dictionary when enabled. public RedBlackTree(IEnumerable sortedKeys, bool enableNodeLookUp = false, IEqualityComparer equalityComparer = null) : this(enableNodeLookUp, equalityComparer) { ValidateCollection(sortedKeys); var nodes = sortedKeys.Select(x => new RedBlackTreeNode(null, x)).ToArray(); Root = (RedBlackTreeNode)ToBST(nodes); assignColors(Root); assignCount(Root); } /// /// Time complexity: O(log(n)) /// public bool HasItem(T value) { if (Root == null) { return false; } if (nodeLookUp != null) { return nodeLookUp.ContainsKey(value); } return find(value) != null; } /// /// Time complexity: O(1) /// internal void Clear() { Root = null; } /// /// Time complexity: O(log(n)) /// public T Max() { var max = Root.FindMax(); return max == null ? default(T) : max.Value; } private RedBlackTreeNode findMax(RedBlackTreeNode node) { return node.FindMax() as RedBlackTreeNode; } /// /// Time complexity: O(log(n)) /// public T Min() { var min = Root.FindMin(); return min == null ? default(T) : min.Value; } /// /// Time complexity: O(log(n)) /// public int IndexOf(T item) { return Root.Position(item); } /// /// Time complexity: O(log(n)) /// public T ElementAt(int index) { if (index < 0 || index >= Count) { throw new ArgumentNullException("index"); } return Root.KthSmallest(index).Value; } //O(log(n)) worst O(n) for unbalanced tree internal RedBlackTreeNode FindNode(T value) { return Root == null ? null : find(value); } //O(log(n)) worst O(n) for unbalanced tree internal bool Exists(T value) { return FindNode(value) != null; } //find the node with the given identifier among descendants of parent and parent //uses pre-order traversal //O(log(n)) worst O(n) for unbalanced tree private RedBlackTreeNode find(T value) { if (nodeLookUp != null) { return nodeLookUp[value] as RedBlackTreeNode; } return Root.Find(value) as RedBlackTreeNode; } /// /// Time complexity: O(log(n)). /// Returns the position (index) of the value in sorted order of this BST. /// public int Insert(T value) { var node = InsertAndReturnNode(value); return node.Position; } /// /// Time complexity: O(log(n)) /// internal RedBlackTreeNode InsertAndReturnNode(T value) { //empty tree if (Root == null) { Root = new RedBlackTreeNode(null, value) { NodeColor = RedBlackTreeNodeColor.Black }; if (nodeLookUp != null) { nodeLookUp[value] = Root; } return Root; } var newNode = insert(Root, value); if (nodeLookUp != null) { nodeLookUp[value] = newNode; } return newNode; } //O(log(n)) always private RedBlackTreeNode insert(RedBlackTreeNode currentNode, T newNodeValue) { while (true) { var compareResult = currentNode.Value.CompareTo(newNodeValue); //current node is less than new item if (compareResult < 0) { //no right child if (currentNode.Right == null) { //insert var node = currentNode.Right = new RedBlackTreeNode(currentNode, newNodeValue); balanceInsertion(currentNode.Right); return node; } currentNode = currentNode.Right; } //current node is greater than new node else if (compareResult > 0) { if (currentNode.Left == null) { //insert var node = currentNode.Left = new RedBlackTreeNode(currentNode, newNodeValue); balanceInsertion(currentNode.Left); return node; } currentNode = currentNode.Left; } else { //duplicate throw new Exception("Item with same key exists"); } } } private void balanceInsertion(RedBlackTreeNode nodeToBalance) { while (true) { if (nodeToBalance == Root) { nodeToBalance.NodeColor = RedBlackTreeNodeColor.Black; break; } //if node to balance is red if (nodeToBalance.NodeColor == RedBlackTreeNodeColor.Red) { //red-red relation; fix it! if (nodeToBalance.Parent.NodeColor == RedBlackTreeNodeColor.Red) { //red sibling if (nodeToBalance.Parent.Sibling != null && nodeToBalance.Parent.Sibling.NodeColor == RedBlackTreeNodeColor.Red) { //mark both children of parent as black and move up balancing nodeToBalance.Parent.Sibling.NodeColor = RedBlackTreeNodeColor.Black; nodeToBalance.Parent.NodeColor = RedBlackTreeNodeColor.Black; //root is always black if (nodeToBalance.Parent.Parent != Root) { nodeToBalance.Parent.Parent.NodeColor = RedBlackTreeNodeColor.Red; } nodeToBalance.UpdateCounts(); nodeToBalance.Parent.UpdateCounts(); nodeToBalance = nodeToBalance.Parent.Parent; } //absent sibling or black sibling else if (nodeToBalance.Parent.Sibling == null || nodeToBalance.Parent.Sibling.NodeColor == RedBlackTreeNodeColor.Black) { if (nodeToBalance.IsLeftChild && nodeToBalance.Parent.IsLeftChild) { var newRoot = nodeToBalance.Parent; swapColors(nodeToBalance.Parent, nodeToBalance.Parent.Parent); rightRotate(nodeToBalance.Parent.Parent); if (newRoot == Root) { Root.NodeColor = RedBlackTreeNodeColor.Black; } nodeToBalance.UpdateCounts(); nodeToBalance = newRoot; } else if (nodeToBalance.IsLeftChild && nodeToBalance.Parent.IsRightChild) { rightRotate(nodeToBalance.Parent); var newRoot = nodeToBalance; swapColors(nodeToBalance.Parent, nodeToBalance); leftRotate(nodeToBalance.Parent); if (newRoot == Root) { Root.NodeColor = RedBlackTreeNodeColor.Black; } nodeToBalance.UpdateCounts(); nodeToBalance = newRoot; } else if (nodeToBalance.IsRightChild && nodeToBalance.Parent.IsRightChild) { var newRoot = nodeToBalance.Parent; swapColors(nodeToBalance.Parent, nodeToBalance.Parent.Parent); leftRotate(nodeToBalance.Parent.Parent); if (newRoot == Root) { Root.NodeColor = RedBlackTreeNodeColor.Black; } nodeToBalance.UpdateCounts(); nodeToBalance = newRoot; } else if (nodeToBalance.IsRightChild && nodeToBalance.Parent.IsLeftChild) { leftRotate(nodeToBalance.Parent); var newRoot = nodeToBalance; swapColors(nodeToBalance.Parent, nodeToBalance); rightRotate(nodeToBalance.Parent); if (newRoot == Root) { Root.NodeColor = RedBlackTreeNodeColor.Black; } nodeToBalance.UpdateCounts(); nodeToBalance = newRoot; } } } } if (nodeToBalance.Parent != null) { nodeToBalance.UpdateCounts(); nodeToBalance = nodeToBalance.Parent; continue; } break; } nodeToBalance.UpdateCounts(true); } private void swapColors(RedBlackTreeNode node1, RedBlackTreeNode node2) { var tmpColor = node2.NodeColor; node2.NodeColor = node1.NodeColor; node1.NodeColor = tmpColor; } /// /// Delete if value exists. /// Time complexity: O(log(n)) /// Returns the position (index) of the item if deleted; otherwise returns -1 /// public int Delete(T value) { if (Root == null) { return -1; } var node = find(value); if (node == null) { return -1; } var position = node.Position; delete(node); if (nodeLookUp != null) { nodeLookUp.Remove(value); } return position; } /// /// Time complexity: O(log(n)) /// public T RemoveAt(int index) { if (index < 0 || index >= Count) { throw new ArgumentNullException("index"); } var node = Root.KthSmallest(index) as RedBlackTreeNode; delete(node); if (nodeLookUp != null) { nodeLookUp.Remove(node.Value); } return node.Value; } //O(log(n)) always private void delete(RedBlackTreeNode node) { //node is a leaf node if (node.IsLeaf) { //if color is red, we are good; no need to balance if (node.NodeColor == RedBlackTreeNodeColor.Red) { deleteLeaf(node); node.Parent?.UpdateCounts(true); return; } deleteLeaf(node); balanceNode(node.Parent); } else { //case one - right tree is null (move sub tree up) if (node.Left != null && node.Right == null) { deleteLeftNode(node); balanceNode(node.Left); } //case two - left tree is null (move sub tree up) else if (node.Right != null && node.Left == null) { deleteRightNode(node); balanceNode(node.Right); } //case three - two child trees //replace the node value with maximum element of left subtree (left max node) //and then delete the left max node else { var maxLeftNode = findMax(node.Left); node.Value = maxLeftNode.Value; if (nodeLookUp != null) { nodeLookUp[node.Value] = node; } //delete left max node delete(maxLeftNode); return; } } } private void balanceNode(RedBlackTreeNode nodeToBalance) { //handle six cases while (nodeToBalance != null) { nodeToBalance.UpdateCounts(); nodeToBalance = handleDoubleBlack(nodeToBalance); } } private void deleteLeaf(RedBlackTreeNode node) { //if node is root if (node.Parent == null) { Root = null; } //assign nodes parent.left/right to null else if (node.IsLeftChild) { node.Parent.Left = null; } else { node.Parent.Right = null; } } private void deleteRightNode(RedBlackTreeNode node) { //root if (node.Parent == null) { Root.Right.Parent = null; Root = Root.Right; Root.NodeColor = RedBlackTreeNodeColor.Black; return; } //node is left child of parent if (node.IsLeftChild) { node.Parent.Left = node.Right; } //node is right child of parent else { node.Parent.Right = node.Right; } node.Right.Parent = node.Parent; if (node.Right.NodeColor != RedBlackTreeNodeColor.Red) { return; } //black deletion! But we can take its red child and recolor it to black //and we are done! node.Right.NodeColor = RedBlackTreeNodeColor.Black; } private void deleteLeftNode(RedBlackTreeNode node) { //root if (node.Parent == null) { Root.Left.Parent = null; Root = Root.Left; Root.NodeColor = RedBlackTreeNodeColor.Black; return; } //node is left child of parent if (node.IsLeftChild) { node.Parent.Left = node.Left; } //node is right child of parent else { node.Parent.Right = node.Left; } node.Left.Parent = node.Parent; if (node.Left.NodeColor != RedBlackTreeNodeColor.Red) { return; } //black deletion! But we can take its red child and recolor it to black //and we are done! node.Left.NodeColor = RedBlackTreeNodeColor.Black; } private void rightRotate(RedBlackTreeNode node) { var prevRoot = node; var leftRightChild = prevRoot.Left.Right; var newRoot = node.Left; //make left child as root prevRoot.Left.Parent = prevRoot.Parent; if (prevRoot.Parent != null) { if (prevRoot.Parent.Left == prevRoot) { prevRoot.Parent.Left = prevRoot.Left; } else { prevRoot.Parent.Right = prevRoot.Left; } } //move prev root as right child of current root newRoot.Right = prevRoot; prevRoot.Parent = newRoot; //move right child of left child of prev root to left child of right child of new root newRoot.Right.Left = leftRightChild; if (newRoot.Right.Left != null) { newRoot.Right.Left.Parent = newRoot.Right; } if (prevRoot == Root) { Root = newRoot; } newRoot.Left.UpdateCounts(); newRoot.Right.UpdateCounts(); newRoot.UpdateCounts(); } private void leftRotate(RedBlackTreeNode node) { var prevRoot = node; var rightLeftChild = prevRoot.Right.Left; var newRoot = node.Right; //make right child as root prevRoot.Right.Parent = prevRoot.Parent; if (prevRoot.Parent != null) { if (prevRoot.Parent.Left == prevRoot) { prevRoot.Parent.Left = prevRoot.Right; } else { prevRoot.Parent.Right = prevRoot.Right; } } //move prev root as left child of current root newRoot.Left = prevRoot; prevRoot.Parent = newRoot; //move left child of right child of prev root to right child of left child of new root newRoot.Left.Right = rightLeftChild; if (newRoot.Left.Right != null) { newRoot.Left.Right.Parent = newRoot.Left; } if (prevRoot == Root) { Root = newRoot; } newRoot.Left.UpdateCounts(); newRoot.Right.UpdateCounts(); newRoot.UpdateCounts(); } private RedBlackTreeNode handleDoubleBlack(RedBlackTreeNode node) { //case 1 if (node == Root) { node.NodeColor = RedBlackTreeNodeColor.Black; return null; } //case 2 if (node.Parent != null && node.Parent.NodeColor == RedBlackTreeNodeColor.Black && node.Sibling != null && node.Sibling.NodeColor == RedBlackTreeNodeColor.Red && ((node.Sibling.Left == null && node.Sibling.Right == null) || (node.Sibling.Left != null && node.Sibling.Right != null && node.Sibling.Left.NodeColor == RedBlackTreeNodeColor.Black && node.Sibling.Right.NodeColor == RedBlackTreeNodeColor.Black))) { node.Parent.NodeColor = RedBlackTreeNodeColor.Red; node.Sibling.NodeColor = RedBlackTreeNodeColor.Black; if (node.Sibling.IsRightChild) { leftRotate(node.Parent); } else { rightRotate(node.Parent); } return node; } //case 3 if (node.Parent != null && node.Parent.NodeColor == RedBlackTreeNodeColor.Black && node.Sibling != null && node.Sibling.NodeColor == RedBlackTreeNodeColor.Black && (node.Sibling.Left == null && node.Sibling.Right == null || node.Sibling.Left != null && node.Sibling.Right != null && node.Sibling.Left.NodeColor == RedBlackTreeNodeColor.Black && node.Sibling.Right.NodeColor == RedBlackTreeNodeColor.Black)) { //pushed up the double black problem up to parent //so now it needs to be fixed node.Sibling.NodeColor = RedBlackTreeNodeColor.Red; return node.Parent; } //case 4 if (node.Parent != null && node.Parent.NodeColor == RedBlackTreeNodeColor.Red && node.Sibling != null && node.Sibling.NodeColor == RedBlackTreeNodeColor.Black && (node.Sibling.Left == null && node.Sibling.Right == null || node.Sibling.Left != null && node.Sibling.Right != null && node.Sibling.Left.NodeColor == RedBlackTreeNodeColor.Black && node.Sibling.Right.NodeColor == RedBlackTreeNodeColor.Black)) { //just swap the color of parent and sibling //which will compensate the loss of black count node.Parent.NodeColor = RedBlackTreeNodeColor.Black; node.Sibling.NodeColor = RedBlackTreeNodeColor.Red; node.UpdateCounts(true); return null; } //case 5 if (node.Parent != null && node.Parent.NodeColor == RedBlackTreeNodeColor.Black && node.Sibling != null && node.Sibling.IsRightChild && node.Sibling.NodeColor == RedBlackTreeNodeColor.Black && node.Sibling.Left != null && node.Sibling.Left.NodeColor == RedBlackTreeNodeColor.Red && node.Sibling.Right != null && node.Sibling.Right.NodeColor == RedBlackTreeNodeColor.Black) { node.Sibling.NodeColor = RedBlackTreeNodeColor.Red; node.Sibling.Left.NodeColor = RedBlackTreeNodeColor.Black; rightRotate(node.Sibling); return node; } //case 5 mirror if (node.Parent != null && node.Parent.NodeColor == RedBlackTreeNodeColor.Black && node.Sibling != null && node.Sibling.IsLeftChild && node.Sibling.NodeColor == RedBlackTreeNodeColor.Black && node.Sibling.Left != null && node.Sibling.Left.NodeColor == RedBlackTreeNodeColor.Black && node.Sibling.Right != null && node.Sibling.Right.NodeColor == RedBlackTreeNodeColor.Red) { node.Sibling.NodeColor = RedBlackTreeNodeColor.Red; node.Sibling.Right.NodeColor = RedBlackTreeNodeColor.Black; leftRotate(node.Sibling); return node; } //case 6 if (node.Parent != null && node.Parent.NodeColor == RedBlackTreeNodeColor.Black && node.Sibling != null && node.Sibling.IsRightChild && node.Sibling.NodeColor == RedBlackTreeNodeColor.Black && node.Sibling.Right != null && node.Sibling.Right.NodeColor == RedBlackTreeNodeColor.Red) { //left rotate to increase the black count on left side by one //and mark the red right child of sibling to black //to compensate the loss of Black on right side of parent node.Sibling.Right.NodeColor = RedBlackTreeNodeColor.Black; leftRotate(node.Parent); node.UpdateCounts(true); return null; } //case 6 mirror if (node.Parent != null && node.Parent.NodeColor == RedBlackTreeNodeColor.Black && node.Sibling != null && node.Sibling.IsLeftChild && node.Sibling.NodeColor == RedBlackTreeNodeColor.Black && node.Sibling.Left != null && node.Sibling.Left.NodeColor == RedBlackTreeNodeColor.Red) { //right rotate to increase the black count on right side by one //and mark the red left child of sibling to black //to compensate the loss of Black on right side of parent node.Sibling.Left.NodeColor = RedBlackTreeNodeColor.Black; rightRotate(node.Parent); node.UpdateCounts(true); return null; } node.UpdateCounts(true); return null; } //assign valid colors assuming the given tree node and its children are in balanced state. private void assignColors(RedBlackTreeNode current) { if (current == null) { return; } assignColors(current.Left); assignColors(current.Right); if (current.IsLeaf) { current.NodeColor = RedBlackTreeNodeColor.Red; } else { current.NodeColor = RedBlackTreeNodeColor.Black; } } /// /// Get the next lower value to given value in this BST. /// public T NextLower(T value) { var node = FindNode(value); if (node == null) { return default(T); } var next = (node as BSTNodeBase).NextLower(); return next != null ? next.Value : default(T); } /// /// Get the next higher to given value in this BST. /// public T NextHigher(T value) { var node = FindNode(value); if (node == null) { return default(T); } var next = (node as BSTNodeBase).NextHigher(); return next != null ? next.Value : default(T); } internal void Swap(T value1, T value2) { var node1 = find(value1); var node2 = find(value2); if (node1 == null || node2 == null) { throw new Exception("Value1, Value2 or both was not found in this BST."); } var tmp = node1.Value; node1.Value = node2.Value; node2.Value = tmp; if (nodeLookUp != null) { nodeLookUp[node1.Value] = node1; nodeLookUp[node2.Value] = node2; } } /// /// Descending enumerable. /// public IEnumerable AsEnumerableDesc() { return GetEnumeratorDesc().AsEnumerable(); } IEnumerator IEnumerable.GetEnumerator() { return GetEnumerator(); } public IEnumerator GetEnumerator() { return new BSTEnumerator(Root); } public IEnumerator GetEnumeratorDesc() { return new BSTEnumerator(Root, false); } } internal enum RedBlackTreeNodeColor { Black, Red } /// /// Red black tree node /// internal class RedBlackTreeNode : BSTNodeBase where T : IComparable { internal new RedBlackTreeNode Parent { get { return (RedBlackTreeNode)base.Parent; } set { base.Parent = value; } } internal new RedBlackTreeNode Left { get { return (RedBlackTreeNode)base.Left; } set { base.Left = value; } } internal new RedBlackTreeNode Right { get { return (RedBlackTreeNode)base.Right; } set { base.Right = value; } } internal RedBlackTreeNodeColor NodeColor { get; set; } internal RedBlackTreeNode Sibling => Parent.Left == this ? Parent.Right : Parent.Left; internal RedBlackTreeNode(RedBlackTreeNode parent, T value) { Parent = parent; Value = value; NodeColor = RedBlackTreeNodeColor.Red; } } }