RedBlackTree.cs

using System;
using System.Collections.Generic;

namespace DataStructures.RedBlackTree;

/// <summary>
///     A self-balancing bindary tree.
/// </summary>
/// <remarks>
///     A red-black tree is a self-balancing binary search tree (BST) that
///     stores a color with each node. A node's color can either be red or
///     black. Several properties are maintained to ensure the tree remains
///     balanced.
///     <list type="number">
///         <item>
///             <term>A red node does not have a red child.</term>
///         </item>
///         <item>
///             <term>All null nodes are considered black.</term>
///         </item>
///         <item>
///             <term>
///                 Every path from a node to its descendant leaf nodes
///             has the same number of black nodes.
///             </term>
///         </item>
///         <item>
///             <term>(Optional) The root is always black.</term>
///         </item>
///     </list>
///     Red-black trees are generally slightly more unbalanced than an
///     AVL tree, but insertion and deletion is generally faster.
///     See https://en.wikipedia.org/wiki/Red%E2%80%93black_tree for more information.
/// </remarks>
/// <typeparam name="TKey">Type of key for the tree.</typeparam>
public class RedBlackTree<TKey>
{
    /// <summary>
    ///     Gets the number of nodes in the tree.
    /// </summary>
    public int Count { get; private set; }

    /// <summary>
    ///     Comparer to use when comparing key values.
    /// </summary>
    private readonly Comparer<TKey> comparer;

    /// <summary>
    ///     Reference to the root node.
    /// </summary>
    private RedBlackTreeNode<TKey>? root;

    /// <summary>
    ///     Initializes a new instance of the <see cref="RedBlackTree{TKey}"/> class.
    /// </summary>
    public RedBlackTree()
    {
        comparer = Comparer<TKey>.Default;
    }

    /// <summary>
    ///     Initializes a new instance of the <see cref="RedBlackTree{TKey}"/> class
    ///     using the specified comparer.
    /// </summary>
    /// <param name="customComparer">Comparer to use when comparing keys.</param>
    public RedBlackTree(Comparer<TKey> customComparer)
    {
        comparer = customComparer;
    }

    /// <summary>
    ///     Add a single node to the tree.
    /// </summary>
    /// <param name="key">Key value to add.</param>
    public void Add(TKey key)
    {
        if (root is null)
        {
            // Case 3
            // New node is root
            root = new RedBlackTreeNode<TKey>(key, null)
            {
                Color = NodeColor.Black,
            };
            Count++;
            return;
        }

        // Regular binary tree insertion
        var node = Add(root, key);

        // Get which side child was added to
        var childDir = comparer.Compare(node.Key, node.Parent!.Key);

        // Set node to be new node's parent for easier handling
        node = node.Parent;

        // Return tree to valid state
        int addCase;
        do
        {
            addCase = GetAddCase(node);

            switch(addCase)
            {
                case 1:
                    break;
                case 2:
                    var oldParent = node.Parent;
                    node = AddCase2(node);

                    if (node is not null)
                    {
                        childDir = comparer.Compare(oldParent!.Key, oldParent.Parent!.Key);
                    }

                    break;
                case 4:
                    node.Color = NodeColor.Black;
                    break;
                case 56:
                    AddCase56(node, comparer.Compare(node.Key, node.Parent!.Key), childDir);
                    break;
                default:
                    throw new InvalidOperationException("It should not be possible to get here!");
            }
        }
        while (addCase == 2 && node is not null);

        Count++;
    }

    /// <summary>
    ///     Add multiple nodes to the tree.
    /// </summary>
    /// <param name="keys">Key values to add.</param>
    public void AddRange(IEnumerable<TKey> keys)
    {
        foreach (var key in keys)
        {
            Add(key);
        }
    }

    /// <summary>
    ///     Remove a node from the tree.
    /// </summary>
    /// <param name="key">Key value to remove.</param>
    public void Remove(TKey key)
    {
        // Search for node
        var node = Remove(root, key);

        // Simple cases
        node = RemoveSimpleCases(node);

        // Exit if deleted node was not non-root black leaf
        if (node is null)
        {
            return;
        }

        // Delete node
        DeleteLeaf(node.Parent!, comparer.Compare(node.Key, node.Parent!.Key));

        // Recolor tree
        do
        {
            node = RemoveRecolor(node);
        }
        while (node is not null && node.Parent is not null);    // Case 2: Reached root

        Count--;
    }

    /// <summary>
    ///     Check if given node is in the tree.
    /// </summary>
    /// <param name="key">Key value to search for.</param>
    /// <returns>Whether or not the node is in the tree.</returns>
    public bool Contains(TKey key)
    {
        var node = root;
        while (node is not null)
        {
            var compareResult = comparer.Compare(key, node.Key);
            if (compareResult < 0)
            {
                node = node.Left;
            }
            else if (compareResult > 0)
            {
                node = node.Right;
            }
            else
            {
                return true;
            }
        }

        return false;
    }

    /// <summary>
    ///     Get the minimum value in the tree.
    /// </summary>
    /// <returns>Minimum value in tree.</returns>
    public TKey GetMin()
    {
        if (root is null)
        {
            throw new InvalidOperationException("Tree is empty!");
        }

        return GetMin(root).Key;
    }

    /// <summary>
    ///     Get the maximum value in the tree.
    /// </summary>
    /// <returns>Maximum value in tree.</returns>
    public TKey GetMax()
    {
        if (root is null)
        {
            throw new InvalidOperationException("Tree is empty!");
        }

        return GetMax(root).Key;
    }

    /// <summary>
    ///     Get keys in order from smallest to largest as defined by the comparer.
    /// </summary>
    /// <returns>Keys in tree in order from smallest to largest.</returns>
    public IEnumerable<TKey> GetKeysInOrder()
    {
        var result = new List<TKey>();
        InOrderWalk(root);
        return result;

        void InOrderWalk(RedBlackTreeNode<TKey>? node)
        {
            if (node is null)
            {
                return;
            }

            InOrderWalk(node.Left);
            result.Add(node.Key);
            InOrderWalk(node.Right);
        }
    }

    /// <summary>
    ///     Get keys in the pre-order order.
    /// </summary>
    /// <returns>Keys in pre-order order.</returns>
    public IEnumerable<TKey> GetKeysPreOrder()
    {
        var result = new List<TKey>();
        PreOrderWalk(root);
        return result;

        void PreOrderWalk(RedBlackTreeNode<TKey>? node)
        {
            if (node is null)
            {
                return;
            }

            result.Add(node.Key);
            PreOrderWalk(node.Left);
            PreOrderWalk(node.Right);
        }
    }

    /// <summary>
    ///     Get keys in the post-order order.
    /// </summary>
    /// <returns>Keys in the post-order order.</returns>
    public IEnumerable<TKey> GetKeysPostOrder()
    {
        var result = new List<TKey>();
        PostOrderWalk(root);
        return result;

        void PostOrderWalk(RedBlackTreeNode<TKey>? node)
        {
            if (node is null)
            {
                return;
            }

            PostOrderWalk(node.Left);
            PostOrderWalk(node.Right);
            result.Add(node.Key);
        }
    }

    /// <summary>
    ///     Perform binary tree insertion.
    /// </summary>
    /// <param name="node">Root of subtree to search from.</param>
    /// <param name="key">Key value to insert.</param>
    /// <returns>Node that was added.</returns>
    private RedBlackTreeNode<TKey> Add(RedBlackTreeNode<TKey> node, TKey key)
    {
        int compareResult;
        RedBlackTreeNode<TKey> newNode;
        while (true)
        {
            compareResult = comparer.Compare(key, node!.Key);
            if (compareResult < 0)
            {
                if (node.Left is null)
                {
                    newNode = new RedBlackTreeNode<TKey>(key, node);
                    node.Left = newNode;
                    break;
                }
                else
                {
                    node = node.Left;
                }
            }
            else if (compareResult > 0)
            {
                if (node.Right is null)
                {
                    newNode = new RedBlackTreeNode<TKey>(key, node);
                    node.Right = newNode;
                    break;
                }
                else
                {
                    node = node.Right;
                }
            }
            else
            {
                throw new ArgumentException($"Key \"{key}\" already exists in tree!");
            }
        }

        return newNode;
    }

    /// <summary>
    ///     Perform case 2 of insertion by pushing blackness down from parent.
    /// </summary>
    /// <param name="node">Parent of inserted node.</param>
    /// <returns>Grandparent of inserted node.</returns>
    private RedBlackTreeNode<TKey>? AddCase2(RedBlackTreeNode<TKey> node)
    {
        var grandparent = node.Parent;
        var parentDir = comparer.Compare(node.Key, node.Parent!.Key);
        var uncle = parentDir < 0 ? grandparent!.Right : grandparent!.Left;

        node.Color = NodeColor.Black;
        uncle!.Color = NodeColor.Black;
        grandparent.Color = NodeColor.Red;

        // Keep root black
        if (node.Parent.Parent is null)
        {
            node.Parent.Color = NodeColor.Black;
        }

        // Set current node as parent to move up tree
        return node.Parent.Parent;
    }

    /// <summary>
    ///     Perform rotations needed for cases 5 and 6 of insertion.
    /// </summary>
    /// <param name="node">Parent of node just inserted.</param>
    /// <param name="parentDir">The side node is on of its parent.</param>
    /// <param name="childDir">The side the child node is on.</param>
    private void AddCase56(RedBlackTreeNode<TKey> node, int parentDir, int childDir)
    {
        if (parentDir < 0)
        {
            // Case 5
            if (childDir > 0)
            {
                node = RotateLeft(node);
            }

            // Case 6
            node = RotateRight(node.Parent!);
            node.Color = NodeColor.Black;
            node.Right!.Color = NodeColor.Red;
        }
        else
        {
            // Case 5
            if (childDir < 0)
            {
                node = RotateRight(node);
            }

            // Case 6
            node = RotateLeft(node.Parent!);
            node.Color = NodeColor.Black;
            node.Left!.Color = NodeColor.Red;
        }
    }

    /// <summary>
    ///     Determine which add case applies to inserted node.
    /// </summary>
    /// <param name="node">Parent of inserted node.</param>
    /// <returns>Case number needed to get tree in valid state. Cases 5 and 6 are represented by 56.</returns>
    private int GetAddCase(RedBlackTreeNode<TKey> node)
    {
        if (node.Color == NodeColor.Black)
        {
            return 1;
        }
        else if (node.Parent is null)
        {
            return 4;
        }
        else
        {
            // Remaining insert cases need uncle
            var grandparent = node.Parent;
            var parentDir = comparer.Compare(node.Key, node.Parent.Key);
            var uncle = parentDir < 0 ? grandparent.Right : grandparent.Left;

            // Case 5 & 6
            if (uncle is null || uncle.Color == NodeColor.Black)
            {
                return 56;
            }

            return 2;
        }
    }

    /// <summary>
    ///     Search for the node to be deleted.
    /// </summary>
    /// <param name="node">Node to start search from.</param>
    /// <param name="key">Key to search for.</param>
    /// <returns>Node to be deleted.</returns>
    private RedBlackTreeNode<TKey> Remove(RedBlackTreeNode<TKey>? node, TKey key)
    {
        if (node is null)
        {
            throw new InvalidOperationException("Tree is empty!");
        }
        else if (!Contains(key))
        {
            throw new KeyNotFoundException($"Key {key} is not in the tree!");
        }
        else
        {
            // Find node
            int dir;
            while (true)
            {
                dir = comparer.Compare(key, node!.Key);
                if (dir < 0)
                {
                    node = node.Left;
                }
                else if (dir > 0)
                {
                    node = node.Right;
                }
                else
                {
                    break;
                }
            }

            return node;
        }
    }

    /// <summary>
    ///     Get the tree back into a valid state after removing non-root black leaf.
    /// </summary>
    /// <param name="node">Non-root black leaf being removed.</param>
    private RedBlackTreeNode<TKey>? RemoveRecolor(RedBlackTreeNode<TKey> node)
    {
        var removeCase = GetRemoveCase(node);

        var dir = comparer.Compare(node.Key, node.Parent!.Key);

        // Determine current node's sibling and nephews
        var sibling = dir < 0 ? node.Parent.Right : node.Parent.Left;
        var closeNewphew = dir < 0 ? sibling!.Left : sibling!.Right;
        var distantNephew = dir < 0 ? sibling!.Right : sibling!.Left;

        switch (removeCase)
        {
            case 1:
                sibling.Color = NodeColor.Red;
                return node.Parent;
            case 3:
                RemoveCase3(node, closeNewphew, dir);
                break;
            case 4:
                RemoveCase4(sibling);
                break;
            case 5:
                RemoveCase5(node, sibling, dir);
                break;
            case 6:
                RemoveCase6(node, distantNephew!, dir);
                break;
            default:
                throw new InvalidOperationException("It should not be possible to get here!");
        }

        return null;
    }

    /// <summary>
    ///     Simple removal cases where black height doesn't change.
    /// </summary>
    /// <param name="node">Node to remove.</param>
    /// <returns>Non-root black leaf node or null. Null indicates that removal was performed.</returns>
    private RedBlackTreeNode<TKey>? RemoveSimpleCases(RedBlackTreeNode<TKey> node)
    {
        // Node to delete is root and has no children
        if (node.Parent is null && node.Left is null && node.Right is null)
        {
            root = null;
            Count--;
            return null;
        }

        // Node has two children. Swap pointers
        if (node.Left is not null && node.Right is not null)
        {
            var successor = GetMin(node.Right);
            node.Key = successor.Key;
            node = successor;
        }

        // At this point node should have at most one child
        if (node.Color == NodeColor.Red)
        {
            // Node is red so it must have no children since it doesn't have two children
            DeleteLeaf(node.Parent!, comparer.Compare(node.Key, node.Parent!.Key));

            Count--;
            return null;
        }
        else
        {
            // Node is black and may or may not be node
            return RemoveBlackNode(node);
        }
    }

    /// <summary>
    ///     Node to delete is black. If it is a leaf then we need to recolor, otherwise remove it.
    /// </summary>
    /// <param name="node">Black node to examine.</param>
    /// <returns>Node to start recoloring from. Null if deletion occurred.</returns>
    private RedBlackTreeNode<TKey>? RemoveBlackNode(RedBlackTreeNode<TKey> node)
    {
        // Node is black and has at most one child. If it has a child it must be red.
        var child = node.Left ?? node.Right;

        // Continue to recoloring if node is leaf
        if (child is null)
        {
            return node;
        }

        // Recolor child
        child.Color = NodeColor.Black;
        child.Parent = node.Parent;

        var childDir = node.Parent is null ? 0 : comparer.Compare(node.Key, node.Parent.Key);

        // Replace node with child
        Transplant(node.Parent, child, childDir);

        Count--;
        return null;
    }

    /// <summary>
    ///     Perform case 3 of removal.
    /// </summary>
    /// <param name="node">Node that was removed.</param>
    /// <param name="closeNephew">Close nephew of removed node.</param>
    /// <param name="childDir">Side of parent the removed node was.</param>
    private void RemoveCase3(RedBlackTreeNode<TKey> node, RedBlackTreeNode<TKey>? closeNephew, int childDir)
    {
        // Rotate and recolor
        var sibling = childDir < 0 ? RotateLeft(node.Parent!) : RotateRight(node.Parent!);
        sibling.Color = NodeColor.Black;
        if (childDir < 0)
        {
            sibling.Left!.Color = NodeColor.Red;
        }
        else
        {
            sibling.Right!.Color = NodeColor.Red;
        }

        // Get new distant newphew
        sibling = closeNephew!;
        var distantNephew = childDir < 0 ? sibling.Right : sibling.Left;

        // Parent is red, sibling is black
        if (distantNephew is not null && distantNephew.Color == NodeColor.Red)
        {
            RemoveCase6(node, distantNephew, childDir);
            return;
        }

        // Get new close nephew
        closeNephew = childDir < 0 ? sibling!.Left : sibling!.Right;

        // Sibling is black, distant nephew is black
        if (closeNephew is not null && closeNephew.Color == NodeColor.Red)
        {
            RemoveCase5(node, sibling!, childDir);
            return;
        }

        // Final recoloring
        RemoveCase4(sibling!);
    }

    /// <summary>
    ///     Perform case 4 of removal.
    /// </summary>
    /// <param name="sibling">Sibling of removed node.</param>
    private void RemoveCase4(RedBlackTreeNode<TKey> sibling)
    {
        sibling.Color = NodeColor.Red;
        sibling.Parent!.Color = NodeColor.Black;
    }

    /// <summary>
    ///     Perform case 5 of removal.
    /// </summary>
    /// <param name="node">Node that was removed.</param>
    /// <param name="sibling">Sibling of removed node.</param>
    /// <param name="childDir">Side of parent removed node was on.</param>
    private void RemoveCase5(RedBlackTreeNode<TKey> node, RedBlackTreeNode<TKey> sibling, int childDir)
    {
        sibling = childDir < 0 ? RotateRight(sibling) : RotateLeft(sibling);
        var distantNephew = childDir < 0 ? sibling.Right! : sibling.Left!;

        sibling.Color = NodeColor.Black;
        distantNephew.Color = NodeColor.Red;

        RemoveCase6(node, distantNephew, childDir);
    }

    /// <summary>
    ///     Perform case 6 of removal.
    /// </summary>
    /// <param name="node">Node that was removed.</param>
    /// <param name="distantNephew">Distant nephew of removed node.</param>
    /// <param name="childDir">Side of parent removed node was on.</param>
    private void RemoveCase6(RedBlackTreeNode<TKey> node, RedBlackTreeNode<TKey> distantNephew, int childDir)
    {
        var oldParent = node.Parent!;
        node = childDir < 0 ? RotateLeft(oldParent) : RotateRight(oldParent);
        node.Color = oldParent.Color;
        oldParent.Color = NodeColor.Black;
        distantNephew.Color = NodeColor.Black;
    }

    /// <summary>
    ///     Determine which removal case is required.
    /// </summary>
    /// <param name="node">Node being removed.</param>
    /// <returns>Which removal case should be performed.</returns>
    private int GetRemoveCase(RedBlackTreeNode<TKey> node)
    {
        var dir = comparer.Compare(node.Key, node.Parent!.Key);

        // Determine current node's sibling and nephews
        var sibling = dir < 0 ? node.Parent.Right : node.Parent.Left;
        var closeNewphew = dir < 0 ? sibling!.Left : sibling!.Right;
        var distantNephew = dir < 0 ? sibling!.Right : sibling!.Left;

        if (sibling.Color == NodeColor.Red)
        {
            return 3;
        }
        else if (distantNephew is not null && distantNephew.Color == NodeColor.Red)
        {
            return 6;
        }
        else if (closeNewphew is not null && closeNewphew.Color == NodeColor.Red)
        {
            return 5;
        }
        else if (node.Parent.Color == NodeColor.Red)
        {
            return 4;
        }
        else
        {
            return 1;
        }
    }

    /// <summary>
    ///     Set child of node or delete leaf.
    /// </summary>
    /// <param name="node">Node to set child of. Set to null for root.</param>
    /// <param name="child">Node to set as child.</param>
    /// <param name="dir">Which side of node to place child.</param>
    private void Transplant(RedBlackTreeNode<TKey>? node, RedBlackTreeNode<TKey>? child, int dir)
    {
        if (node is null)
        {
            root = child;
        }
        else if (child is null)
        {
            DeleteLeaf(node, dir);
        }
        else if (dir < 0)
        {
            node.Left = child;
        }
        else
        {
            node.Right = child;
        }
    }

    /// <summary>
    ///     Delete leaf node.
    /// </summary>
    /// <param name="node">Parent of leaf node to delete.</param>
    /// <param name="dir">Side of parent leaf is on.</param>
    private void DeleteLeaf(RedBlackTreeNode<TKey> node, int dir)
    {
        if (dir < 0)
        {
            node.Left = null;
        }
        else
        {
            node.Right = null;
        }
    }

    /// <summary>
    ///     Perform a left (counter-clockwise) rotation.
    /// </summary>
    /// <param name="node">Node to rotate about.</param>
    /// <returns>New node with rotation applied.</returns>
    private RedBlackTreeNode<TKey> RotateLeft(RedBlackTreeNode<TKey> node)
    {
        var temp1 = node;
        var temp2 = node!.Right!.Left;

        node = node.Right;
        node.Parent = temp1.Parent;
        if (node.Parent is not null)
        {
            var nodeDir = comparer.Compare(node.Key, node.Parent.Key);
            if (nodeDir < 0)
            {
                node.Parent.Left = node;
            }
            else
            {
                node.Parent.Right = node;
            }
        }

        node.Left = temp1;
        node.Left.Parent = node;

        node.Left.Right = temp2;
        if (temp2 is not null)
        {
            node.Left.Right!.Parent = temp1;
        }

        if (node.Parent is null)
        {
            root = node;
        }

        return node;
    }

    /// <summary>
    ///     Perform a right (clockwise) rotation.
    /// </summary>
    /// <param name="node">Node to rotate about.</param>
    /// <returns>New node with rotation applied.</returns>
    private RedBlackTreeNode<TKey> RotateRight(RedBlackTreeNode<TKey> node)
    {
        var temp1 = node;
        var temp2 = node!.Left!.Right;

        node = node.Left;
        node.Parent = temp1.Parent;
        if (node.Parent is not null)
        {
            var nodeDir = comparer.Compare(node.Key, node.Parent.Key);
            if (nodeDir < 0)
            {
                node.Parent.Left = node;
            }
            else
            {
                node.Parent.Right = node;
            }
        }

        node.Right = temp1;
        node.Right.Parent = node;

        node.Right.Left = temp2;
        if (temp2 is not null)
        {
            node.Right.Left!.Parent = temp1;
        }

        if (node.Parent is null)
        {
            root = node;
        }

        return node;
    }

    /// <summary>
    ///     Helper function to get node instance with minimum key value
    ///     in the specified subtree.
    /// </summary>
    /// <param name="node">Node specifying root of subtree.</param>
    /// <returns>Minimum value in node's subtree.</returns>
    private RedBlackTreeNode<TKey> GetMin(RedBlackTreeNode<TKey> node)
    {
        while (node.Left is not null)
        {
            node = node.Left;
        }

        return node;
    }

    /// <summary>
    ///     Helper function to get node instance with maximum key value
    ///     in the specified subtree.
    /// </summary>
    /// <param name="node">Node specifyng root of subtree.</param>
    /// <returns>Maximum value in node's subtree.</returns>
    private RedBlackTreeNode<TKey> GetMax(RedBlackTreeNode<TKey> node)
    {
        while (node.Right is not null)
        {
            node = node.Right;
        }

        return node;
    }
}