Add a weight-balanced tree

Quadratic.Carto.Craft2/Containers/WBTree.cs

@@ -0,0 +1,164 @@
+using System.Diagnostics.CodeAnalysis;
+using System.Drawing;
+
+namespace Quadratic.Carto.Craft2.Containers;
+
+/// <summary>
+/// Weight-balanced tree. Immutable.
+/// 
+/// Also see: <a href="https://en.wikipedia.org/wiki/Weight-balanced_tree">Wikipedia</a>,
+/// the <a href="https://dl.acm.org/doi/pdf/10.1145/1499949.1500040">original paper</a>,
+/// and the <a href="https://yoichihirai.com/bst.pdf">paper</a> by Hirai about selecting the right balance factor.
+/// </summary>
+public struct WBTree<T> where T : IComparable<T>
+{
+    Node? root;
+
+    class Node(T value, Node? left = null, Node? right = null)
+    {
+        public int Size { get; } = NSize(left) + NSize(right) + 1;
+        public T Value { get; set; } = value;
+        public Node? Left { get; } = left;
+        public Node? Right { get; } = right;
+
+        public void Deconstruct(out T value, out Node? left, out Node? right)
+        {
+            value = this.Value;
+            left = this.Left;
+            right = this.Right;
+        }
+
+        public bool Balanced()
+        {
+            return IsBalanced(Left, Right) && IsBalanced(Right, Left) &&
+                (Left?.Balanced() ?? true) && (Right?.Balanced() ?? true);
+        }
+
+        private static int NSize(Node? n) => n?.Size ?? 0;
+
+        private static bool IsBalanced(Node? l, Node? r) =>
+            DELTA * (NSize(l) + 1) >= (NSize(r) + 1);
+
+        private static bool IsSingle(Node? l, Node? r) =>
+            (NSize(l) + 1) < GAMMA * (NSize(r) + 1);
+
+        private const int DELTA = 2;
+        private const int GAMMA = 3;
+
+        public static Node Insert(Node? node, T value)
+        {
+            if (node == null) return new Node(value);
+            var cmp = value.CompareTo(node.Value);
+            if (cmp < 0)
+                return BalanceR(node.Value, Insert(node.Left, value), node.Right);
+            else if (cmp > 0)
+                return BalanceL(node.Value, node.Left, Insert(node.Right, value));
+            else
+                return new Node(value, node.Left, node.Right);
+        }
+
+        private static Node BalanceL(T v, Node? l, Node r)
+        {
+            if (IsBalanced(l, r)) return new Node(v, l, r);
+            // rotateL
+            else if (IsSingle(l, r))
+            {
+                var (k2, t2, t3) = r;
+                return new Node(k2, new Node(v, l, t2), t3);
+            }
+            else
+            {
+                var (k2, rl, t4) = r;
+                if (rl is null)
+                    throw new ArgumentException("Unable to perform double rotation; r.Left is null");
+                var (k3, t2, t3) = rl;
+                return new Node(k3, new Node(v, l, t2), new Node(k2, t3, t4));
+            }
+        }
+
+        private static Node BalanceR(T cVal, Node l, Node? r)
+        {
+            if (IsBalanced(r, l)) return new Node(cVal, l, r);
+            // rotateR
+            else if (IsSingle(r, l))
+            {
+                var (k2, t2, t3) = l;
+                return new Node(k2, t2, new Node(cVal, t3, r));
+            }
+            else
+            {
+                var (k2, t2, lr) = l;
+                if (lr is null)
+                    throw new ArgumentException("Unable to perform double rotation; l.Right is null");
+                var (k3, t3, t4) = lr;
+                return new Node(k3, new Node(k2, t2, t3), new Node(cVal, t4, r));
+            }
+        }
+    }
+
+    private WBTree(Node? node)
+    {
+        this.root = node;
+    }
+
+    public static WBTree<T> Empty => new WBTree<T>(null);
+
+    public static WBTree<T> Singleton(T value) => new WBTree<T>(new Node(value));
+
+    public readonly int Count => root is null ? 0 : root.Size;
+
+    public WBTree<T> Add(T value) => new WBTree<T>(Node.Insert(root, value));
+
+    public bool TryGet(T value, [NotNullWhen(true)] out T? result)
+    {
+        var node = root;
+        while (node != null)
+        {
+            var cmp = value.CompareTo(node.Value);
+            if (cmp < 0)
+                node = node.Left;
+            else if (cmp > 0)
+                node = node.Right;
+            else
+            {
+                result = node.Value;
+                return true;
+            }
+        }
+        result = default;
+        return false;
+    }
+
+    public bool Contains(T value) => TryGet(value, out _);
+
+    public T Get(T value) => TryGet(value, out var result) ? result : throw new KeyNotFoundException();
+}
+
+/// <summary>
+/// A mutable reference to a weight-balanced tree. Might be preferable to use in some cases.
+/// </summary>
+/// <typeparam name="T"></typeparam>
+public class WBTreeRef<T> where T : IComparable<T>
+{
+    WBTree<T> impl;
+
+    public WBTreeRef() => impl = WBTree<T>.Empty;
+
+    public WBTreeRef(T value) => impl = WBTree<T>.Singleton(value);
+
+    public int Count => impl.Count;
+
+    public void Add(T value) => impl = impl.Add(value);
+
+    public bool TryGet(T value, [NotNullWhen(true)] out T? result) => impl.TryGet(value, out result);
+
+    public bool Contains(T value) => impl.Contains(value);
+
+    public T Get(T value) => impl.Get(value);
+
+    public WBTree<T> ToImmutable() => impl;
+
+    public static implicit operator WBTreeRef<T>(WBTree<T> tree) => new() { impl = tree };
+
+    public static implicit operator WBTree<T>(WBTreeRef<T> tree) => tree.impl;
+}