Balanced BSTs & ordered sets

        y                x
       / \              / \
      x   C    ==>     A   y
     / \                  / \
    A   B                B   C

package main
 
import (
	"cmp"
	"fmt"
	"math/rand/v2"
)
 
type node[K cmp.Ordered] struct {
	key         K
	prio        uint64
	size        int
	left, right *node[K]
}
 
func size[K cmp.Ordered](n *node[K]) int {
	if n == nil {
		return 0
	}
	return n.size
}
 
func (n *node[K]) fix() { n.size = 1 + size(n.left) + size(n.right) }
 
func rotateRight[K cmp.Ordered](y *node[K]) *node[K] {
	x := y.left
	y.left, x.right = x.right, y
	y.fix()
	x.fix()
	return x
}
 
func rotateLeft[K cmp.Ordered](x *node[K]) *node[K] {
	y := x.right
	x.right, y.left = y.left, x
	x.fix()
	y.fix()
	return y
}
 
// Set is an ordered set of keys backed by a treap.
type Set[K cmp.Ordered] struct {
	root *node[K]
	rng  *rand.Rand
}
 
func NewSet[K cmp.Ordered](seed uint64) *Set[K] {
	return &Set[K]{rng: rand.New(rand.NewPCG(seed, seed))}
}
 
func (s *Set[K]) Len() int { return size(s.root) }
 
func (s *Set[K]) Insert(k K) {
	var ins func(n *node[K]) *node[K]
	ins = func(n *node[K]) *node[K] {
		if n == nil {
			return &node[K]{key: k, prio: s.rng.Uint64(), size: 1}
		}
		switch {
		case k < n.key:
			n.left = ins(n.left)
			if n.left.prio > n.prio {
				n = rotateRight(n) // the child has higher priority: lift it
			}
		case k > n.key:
			n.right = ins(n.right)
			if n.right.prio > n.prio {
				n = rotateLeft(n)
			}
		}
		n.fix()
		return n
	}
	s.root = ins(s.root)
}
 
func (s *Set[K]) Delete(k K) {
	var del func(n *node[K]) *node[K]
	del = func(n *node[K]) *node[K] {
		if n == nil {
			return nil
		}
		switch {
		case k < n.key:
			n.left = del(n.left)
		case k > n.key:
			n.right = del(n.right)
		default:
			// Rotate the node down toward a leaf, lifting the higher-priority child, then drop it.
			if n.left == nil {
				return n.right
			}
			if n.right == nil {
				return n.left
			}
			if n.left.prio > n.right.prio {
				n = rotateRight(n)
				n.right = del(n.right)
			} else {
				n = rotateLeft(n)
				n.left = del(n.left)
			}
		}
		n.fix()
		return n
	}
	s.root = del(s.root)
}
 
func (s *Set[K]) Contains(k K) bool {
	for n := s.root; n != nil; {
		switch {
		case k < n.key:
			n = n.left
		case k > n.key:
			n = n.right
		default:
			return true
		}
	}
	return false
}
 
// Rank returns how many keys are smaller than k.
func (s *Set[K]) Rank(k K) int {
	r := 0
	for n := s.root; n != nil; {
		if k <= n.key {
			n = n.left
		} else {
			r += size(n.left) + 1
			n = n.right
		}
	}
	return r
}
 
// Kth returns the k-th smallest key (0-based).
func (s *Set[K]) Kth(k int) K {
	n := s.root
	for {
		switch l := size(n.left); {
		case k < l:
			n = n.left
		case k == l:
			return n.key
		default:
			k -= l + 1
			n = n.right
		}
	}
}
 
// Ceil returns the smallest key >= k.
func (s *Set[K]) Ceil(k K) (K, bool) {
	var best K
	found := false
	for n := s.root; n != nil; {
		if n.key >= k {
			best, found = n.key, true
			n = n.left
		} else {
			n = n.right
		}
	}
	return best, found
}
 
func (s *Set[K]) Keys() []K {
	var out []K
	var walk func(n *node[K])
	walk = func(n *node[K]) {
		if n != nil {
			walk(n.left)
			out = append(out, n.key)
			walk(n.right)
		}
	}
	walk(s.root)
	return out
}
 
func height[K cmp.Ordered](n *node[K]) int {
	if n == nil {
		return 0
	}
	return 1 + max(height(n.left), height(n.right))
}
 
func main() {
	s := NewSet[int](42)
	for _, k := range []int{50, 20, 80, 10, 30, 70, 90, 30} {
		s.Insert(k) // the duplicate 30 is ignored
	}
	fmt.Println("keys:", s.Keys(), "len", s.Len())
	fmt.Println("contains 70:", s.Contains(70), " contains 75:", s.Contains(75))
	fmt.Println("rank(70):", s.Rank(70), " 3rd smallest:", s.Kth(2))
	c, ok := s.Ceil(55)
	fmt.Println("ceil(55):", c, ok)
	s.Delete(50)
	fmt.Println("after deleting 50:", s.Keys())
 
	// The real test: sorted input.
	big := NewSet[int](7)
	for i := range 100_000 {
		big.Insert(i)
	}
	fmt.Println("100000 sorted inserts, height:", height(big.root), "(a plain BST would be 100000)")
	fmt.Println("median:", big.Kth(big.Len()/2), " rank(12345):", big.Rank(12345))
}

keys: [10 20 30 50 70 80 90] len 7
contains 70: true  contains 75: false
rank(70): 4  3rd smallest: 30
ceil(55): 70 true
after deleting 50: [10 20 30 70 80 90]
100000 sorted inserts, height: 41 (a plain BST would be 100000)
median: 50000  rank(12345): 12345