GrokAlgorithms/Chapter8_BalancedTree/AvlTree.cs

using System.Collections;

class AvlTree<T> : IEnumerable<T> where T : notnull
{
    private IComparer<T> _comparer;
    private TreeNode? _root;

    public AvlTree(IComparer<T> comparer)
    {
        _comparer = comparer;
    }

    public void Insert(T value)
    {
        _root = Insert(_root, value, _comparer);
    }

    public void Remove(T value)
    {
        _root = Remove(_root, value, _comparer);
    }

    public int Count(T value)
    {
        var curNode = _root;

        while (curNode != null)
        {
            var cmpResult = _comparer.Compare(value, curNode.Value);
            if (cmpResult == 0)
                return curNode.Count;
            else if (cmpResult > 0)
                curNode = curNode.Right;
            else if (cmpResult < 0)
                curNode = curNode.Left;
        }
        return 0;
    }

    public override string ToString()
    {
        return PrintTree(_root, "", true);

        string PrintTree(TreeNode? node, string indent, bool last)
        {
            if (node == null)
                return "";

            var result = indent;
            result += last ? "└─" : "├─";
            result += $"{node.Value} (h={node.Height}, c={node.Count})\n";

            indent += last ? "  " : "│ ";

            var children = new List<TreeNode>();
            if (node.Right != null) children.Add(node.Right);
            if (node.Left != null) children.Add(node.Left);

            for (int i = 0; i < children.Count; i++)
            {
                result += PrintTree(children[i], indent, i == children.Count - 1);
            }
            return result;
        }
    }

    private static TreeNode Insert(TreeNode? node, T value, IComparer<T> comparer)
    {
        if (node == null)
            return new TreeNode() { Value = value };
        var cmpResult = comparer.Compare(value, node.Value);
        if (cmpResult < 0)
            node.Left = Insert(node.Left, value, comparer);
        else if (cmpResult > 0)
            node.Right = Insert(node.Right, value, comparer);
        else
            node.Count++;
        return Balance(node);
    }

    private static TreeNode? Remove(TreeNode? node, T value, IComparer<T> comparer)
    {
        if (node == null)
            return null;

        var cmpResult = comparer.Compare(value, node.Value);
        if (cmpResult != 0)
        {
            if (cmpResult > 0)
                node.Right = Remove(node.Right, value, comparer);
            else if (cmpResult < 0)
                node.Left = Remove(node.Left, value, comparer);
            return Balance(node);
        }

        node.Count--;
        if (node.Count > 0) // Если после уменьшения счётчика ещё есть элементы
            return node;


        if (node.Right == null) // Простой случай, когда нет правого узла - возвращаем левый
            return node.Left;

        TreeNode minRight;
        var newRight = RemoveMin(node.Right, out minRight);
        minRight.Left = node.Left;
        minRight.Right = newRight;
        return Balance(minRight);
    }

    private static TreeNode? RemoveMin(TreeNode node, out TreeNode minNode)
    {
        if (node.Left == null)
        {
            minNode = node;
            return node.Right;
        }
        node.Left = RemoveMin(node.Left, out minNode);
        return Balance(node);
    }

    private static TreeNode RotateRight(TreeNode p)
    {
        var q = p.Left!;
        p.Left = q.Right;
        q.Right = p;
        p.UpdateHeight();
        q.UpdateHeight();
        return q;
    }

    private static TreeNode RotateLeft(TreeNode p)
    {
        var q = p.Right!;
        p.Right = q.Left;
        q.Left = p;
        p.UpdateHeight();
        q.UpdateHeight();
        return q;
    }

    private static TreeNode Balance(TreeNode node)
    {
        node.UpdateHeight();
        if (node.BFactor > 1)
        {
            if (node.Right?.BFactor < 0)
                node.Right = RotateRight(node.Right);
            return RotateLeft(node);
        }
        else if (node.BFactor < -1)
        {
            if (node.Left?.BFactor > 0)
                node.Left = RotateLeft(node.Left);
            return RotateRight(node);
        }
        return node;
    }

    private IEnumerable<T> Dfs(TreeNode? node)
    {
        if (node == null)
            yield break;

        foreach (var val in Dfs(node.Left))
            yield return val;

        for (int i = 0; i < node.Count; i++)
            yield return node.Value;

        foreach (var val in Dfs(node.Right))
            yield return val;
    }

    public IEnumerable<T> TraverseDfs()
    {
        return Dfs(_root);
    }

    public IEnumerator<T> GetEnumerator()
    {
        return TraverseDfs().GetEnumerator();
    }

    IEnumerator IEnumerable.GetEnumerator()
    {
        return TraverseDfs().GetEnumerator();
    }

    private class TreeNode
    {
        public required T Value;
        public TreeNode? Left;
        public TreeNode? Right;
        public int Height { get; private set; } = 1;
        public int BFactor { get; private set; } = 0;
        public int Count = 1;

        public void UpdateHeight()
        {
            Height = Math.Max(Left?.Height ?? 0, Right?.Height ?? 0) + 1;
            BFactor = (Right?.Height ?? 0) - (Left?.Height ?? 0);
        }
    }
}