AATree.cs

using System;
using System.Collections.Generic;

namespace DataStructures.AATree;

/// <summary>
///     A simple self-balancing binary search tree.
/// </summary>
/// <remarks>
///     AA Trees are a form of self-balancing binary search tree named after their inventor
///     Arne Anderson. AA Trees are designed to be simple to understand and implement.
///     This is accomplished by limiting how nodes can be added to the tree.
///     This simplifies rebalancing operations.
///     More information: https://en.wikipedia.org/wiki/AA_tree .
/// </remarks>
/// <typeparam name="TKey">The type of key for the AA tree.</typeparam>
public class AaTree<TKey>
{
    /// <summary>
    ///     The comparer function to use to compare the keys.
    /// </summary>
    private readonly Comparer<TKey> comparer;

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

    /// <summary>
    ///     Initializes a new instance of the <see cref="AaTree{TKey}" /> class with a custom comparer.
    /// </summary>
    /// <param name="customComparer">The custom comparer to use to compare keys.</param>
    public AaTree(Comparer<TKey> customComparer) => comparer = customComparer;

    /// <summary>
    ///     Gets the root of the tree.
    /// </summary>
    public AaTreeNode<TKey>? Root { get; private set; }

    /// <summary>
    ///     Gets the number of elements in the tree.
    /// </summary>
    public int Count { get; private set; }

    /// <summary>
    ///     Add a single element to the tree.
    /// </summary>
    /// <param name="key">The element to add to the tree.</param>
    public void Add(TKey key)
    {
        Root = Add(key, Root);
        Count++;
    }

    /// <summary>
    ///     Add multiple elements to the tree.
    /// </summary>
    /// <param name="keys">The elements to add to the tree.</param>
    public void AddRange(IEnumerable<TKey> keys)
    {
        foreach (var key in keys)
        {
            Root = Add(key, Root);
            Count++;
        }
    }

    /// <summary>
    ///     Remove a single element from the tree.
    /// </summary>
    /// <param name="key">Element to remove.</param>
    public void Remove(TKey key)
    {
        if (!Contains(key, Root))
        {
            throw new InvalidOperationException($"{nameof(key)} is not in the tree");
        }

        Root = Remove(key, Root);
        Count--;
    }

    /// <summary>
    ///     Checks if the specified element is in the tree.
    /// </summary>
    /// <param name="key">The element to look for.</param>
    /// <returns>true if the element is in the tree, false otherwise.</returns>
    public bool Contains(TKey key) => Contains(key, Root);

    /// <summary>
    ///     Gets the largest element in the tree. (ie. the element in the right most node).
    /// </summary>
    /// <returns>The largest element in the tree according to the stored comparer.</returns>
    /// <exception cref="InvalidOperationException">Thrown if the tree is empty.</exception>
    public TKey GetMax()
    {
        if (Root is null)
        {
            throw new InvalidOperationException("Tree is empty!");
        }

        return GetMax(Root).Key;
    }

    /// <summary>
    ///     Gets the smallest element in the tree. (ie. the element in the left most node).
    /// </summary>
    /// <returns>The smallest element in the tree according to the stored comparer.</returns>
    /// <throws>InvalidOperationException if the tree is empty.</throws>
    public TKey GetMin()
    {
        if (Root is null)
        {
            throw new InvalidOperationException("Tree is empty!");
        }

        return GetMin(Root).Key;
    }

    /// <summary>
    ///     Gets all the elements in the tree in in-order order.
    /// </summary>
    /// <returns>Sequence of elements in in-order order.</returns>
    public IEnumerable<TKey> GetKeysInOrder()
    {
        var result = new List<TKey>();
        InOrderWalk(Root);
        return result;

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

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

    /// <summary>
    ///     Gets all the elements in the tree in pre-order order.
    /// </summary>
    /// <returns>Sequence of elements in pre-order order.</returns>
    public IEnumerable<TKey> GetKeysPreOrder()
    {
        var result = new List<TKey>();
        PreOrderWalk(Root);
        return result;

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

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

    /// <summary>
    ///     Gets all the elements in the tree in post-order order.
    /// </summary>
    /// <returns>Sequence of elements in post-order order.</returns>
    public IEnumerable<TKey> GetKeysPostOrder()
    {
        var result = new List<TKey>();
        PostOrderWalk(Root);
        return result;

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

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

    /// <summary>
    ///     Recursive function to add an element to the tree.
    /// </summary>
    /// <param name="key">The element to add.</param>
    /// <param name="node">The node to search for a empty spot.</param>
    /// <returns>The node with the added element.</returns>
    /// <exception cref="ArgumentException">Thrown if key is already in the tree.</exception>
    private AaTreeNode<TKey> Add(TKey key, AaTreeNode<TKey>? node)
    {
        if (node is null)
        {
            return new AaTreeNode<TKey>(key, 1);
        }

        if (comparer.Compare(key, node.Key) < 0)
        {
            node.Left = Add(key, node.Left);
        }
        else if (comparer.Compare(key, node.Key) > 0)
        {
            node.Right = Add(key, node.Right);
        }
        else
        {
            throw new ArgumentException($"Key \"{key}\" already in tree!", nameof(key));
        }

        return Split(Skew(node))!;
    }

    /// <summary>
    ///     Recursive function to remove an element from the tree.
    /// </summary>
    /// <param name="key">The element to remove.</param>
    /// <param name="node">The node to search from.</param>
    /// <returns>The node with the specified element removed.</returns>
    private AaTreeNode<TKey>? Remove(TKey key, AaTreeNode<TKey>? node)
    {
        if (node is null)
        {
            return null;
        }

        if (comparer.Compare(key, node.Key) < 0)
        {
            node.Left = Remove(key, node.Left);
        }
        else if (comparer.Compare(key, node.Key) > 0)
        {
            node.Right = Remove(key, node.Right);
        }
        else
        {
            if (node.Left is null && node.Right is null)
            {
                return null;
            }

            if (node.Left is null)
            {
                var successor = GetMin(node.Right!);
                node.Right = Remove(successor.Key, node.Right);
                node.Key = successor.Key;
            }
            else
            {
                var predecessor = GetMax(node.Left);
                node.Left = Remove(predecessor.Key, node.Left);
                node.Key = predecessor.Key;
            }
        }

        node = DecreaseLevel(node);
        node = Skew(node);
        node!.Right = Skew(node.Right);

        if (node.Right is not null)
        {
            node.Right.Right = Skew(node.Right.Right);
        }

        node = Split(node);
        node!.Right = Split(node.Right);
        return node;
    }

    /// <summary>
    ///     Recursive function to check if the element exists in the tree.
    /// </summary>
    /// <param name="key">The element to check for.</param>
    /// <param name="node">The node to search from.</param>
    /// <returns>true if the element exists in the tree, false otherwise.</returns>
    private bool Contains(TKey key, AaTreeNode<TKey>? node) =>
        node is { }
        && comparer.Compare(key, node.Key) is { } v
        && v switch
        {
            { } when v > 0 => Contains(key, node.Right),
            { } when v < 0 => Contains(key, node.Left),
            _ => true,
        };

    /// <summary>
    ///     Recursive to find the maximum/right-most element.
    /// </summary>
    /// <param name="node">The node to traverse from.</param>
    /// <returns>The node with the maximum/right-most element.</returns>
    private AaTreeNode<TKey> GetMax(AaTreeNode<TKey> node)
    {
        while (true)
        {
            if (node.Right is null)
            {
                return node;
            }

            node = node.Right;
        }
    }

    /// <summary>
    ///     Recursive to find the minimum/left-most element.
    /// </summary>
    /// <param name="node">The node to traverse from.</param>
    /// <returns>The node with the minimum/left-most element.</returns>
    private AaTreeNode<TKey> GetMin(AaTreeNode<TKey> node)
    {
        while (true)
        {
            if (node.Left is null)
            {
                return node;
            }

            node = node.Left;
        }
    }

    /// <summary>
    ///     Remove right-horizontal links and replaces them with left-horizontal links.
    ///     Accomplishes this by performing a right rotation.
    /// </summary>
    /// <param name="node">The node to rebalance from.</param>
    /// <returns>The rebalanced node.</returns>
    private AaTreeNode<TKey>? Skew(AaTreeNode<TKey>? node)
    {
        if (node?.Left is null || node.Left.Level != node.Level)
        {
            return node;
        }

        var left = node.Left;
        node.Left = left.Right;
        left.Right = node;
        return left;
    }

    /// <summary>
    ///     Reduces the number of right-horizontal links.
    ///     Accomplishes this by performing a left rotation, and incrementing level.
    /// </summary>
    /// <param name="node">The node to rebalance from.</param>
    /// <returns>The rebalanced node.</returns>
    private AaTreeNode<TKey>? Split(AaTreeNode<TKey>? node)
    {
        if (node?.Right?.Right is null || node.Level != node.Right.Right.Level)
        {
            return node;
        }

        var right = node.Right;
        node.Right = right.Left;
        right.Left = node;
        right.Level++;
        return right;
    }

    /// <summary>
    ///     Decreases the level of node if necessary.
    /// </summary>
    /// <param name="node">The node to decrease level from.</param>
    /// <returns>The node with modified level.</returns>
    private AaTreeNode<TKey> DecreaseLevel(AaTreeNode<TKey> node)
    {
        var newLevel = Math.Min(GetLevel(node.Left), GetLevel(node.Right)) + 1;
        if (newLevel >= node.Level)
        {
            return node;
        }

        node.Level = newLevel;
        if (node.Right is { } && newLevel < node.Right.Level)
        {
            node.Right.Level = newLevel;
        }

        return node;

        static int GetLevel(AaTreeNode<TKey>? x) => x?.Level ?? 0;
    }
}