avl.go

// Package avl implements a generic AVL self-balancing binary search tree.
package avl

// Ordered constrains types that support the < operator.
type Ordered interface {
	~int | ~int8 | ~int16 | ~int32 | ~int64 |
		~uint | ~uint8 | ~uint16 | ~uint32 | ~uint64 | ~uintptr |
		~float32 | ~float64 | ~string
}

// node is a single AVL tree node.
type node[T Ordered] struct {
	key         T
	height      int
	left, right *node[T]
}

// Tree is the public AVL tree handle. The zero value is a valid empty tree.
type Tree[T Ordered] struct {
	root *node[T]
	size int
}

// Len returns the number of distinct keys in the tree.
func (t *Tree[T]) Len() int { return t.size }

// Height returns the height of the tree (0 for empty).
func (t *Tree[T]) Height() int { return height(t.root) }

// Insert adds key to the tree. Returns true if it was newly inserted, false
// if it was already present.
func (t *Tree[T]) Insert(key T) bool {
	var inserted bool
	t.root, inserted = insert(t.root, key)
	if inserted {
		t.size++
	}
	return inserted
}

// Contains reports whether key is present.
func (t *Tree[T]) Contains(key T) bool {
	for n := t.root; n != nil; {
		switch {
		case key < n.key:
			n = n.left
		case key > n.key:
			n = n.right
		default:
			return true
		}
	}
	return false
}

// Delete removes key. Returns true if a node was actually removed.
func (t *Tree[T]) Delete(key T) bool {
	var removed bool
	t.root, removed = del(t.root, key)
	if removed {
		t.size--
	}
	return removed
}

// InOrder returns all keys in ascending order.
func (t *Tree[T]) InOrder() []T {
	out := make([]T, 0, t.size)
	var walk func(n *node[T])
	walk = func(n *node[T]) {
		if n == nil {
			return
		}
		walk(n.left)
		out = append(out, n.key)
		walk(n.right)
	}
	walk(t.root)
	return out
}

// --- internal helpers ---

func height[T Ordered](n *node[T]) int {
	if n == nil {
		return 0
	}
	return n.height
}

func updateHeight[T Ordered](n *node[T]) {
	lh, rh := height(n.left), height(n.right)
	if lh > rh {
		n.height = lh + 1
	} else {
		n.height = rh + 1
	}
}

func balance[T Ordered](n *node[T]) int {
	if n == nil {
		return 0
	}
	return height(n.left) - height(n.right)
}

func rotateRight[T Ordered](y *node[T]) *node[T] {
	x := y.left
	t2 := x.right
	x.right = y
	y.left = t2
	updateHeight(y)
	updateHeight(x)
	return x
}

func rotateLeft[T Ordered](x *node[T]) *node[T] {
	y := x.right
	t2 := y.left
	y.left = x
	x.right = t2
	updateHeight(x)
	updateHeight(y)
	return y
}

// rebalance restores the AVL invariant for n and returns the (possibly new)
// root of the subtree.
func rebalance[T Ordered](n *node[T]) *node[T] {
	updateHeight(n)
	b := balance(n)

	switch {
	case b > 1 && balance(n.left) >= 0:
		return rotateRight(n) // LL
	case b > 1 && balance(n.left) < 0:
		n.left = rotateLeft(n.left) // LR
		return rotateRight(n)
	case b < -1 && balance(n.right) <= 0:
		return rotateLeft(n) // RR
	case b < -1 && balance(n.right) > 0:
		n.right = rotateRight(n.right) // RL
		return rotateLeft(n)
	}
	return n
}

func insert[T Ordered](n *node[T], key T) (*node[T], bool) {
	if n == nil {
		return &node[T]{key: key, height: 1}, true
	}
	var inserted bool
	switch {
	case key < n.key:
		n.left, inserted = insert(n.left, key)
	case key > n.key:
		n.right, inserted = insert(n.right, key)
	default:
		return n, false // duplicate
	}
	if !inserted {
		return n, false
	}
	return rebalance(n), true
}

// minNode returns the minimum-key node of the subtree rooted at n.
func minNode[T Ordered](n *node[T]) *node[T] {
	for n.left != nil {
		n = n.left
	}
	return n
}

func del[T Ordered](n *node[T], key T) (*node[T], bool) {
	if n == nil {
		return nil, false
	}
	var removed bool
	switch {
	case key < n.key:
		n.left, removed = del(n.left, key)
	case key > n.key:
		n.right, removed = del(n.right, key)
	default:
		// Found node to delete.
		if n.left == nil || n.right == nil {
			var child *node[T]
			if n.left != nil {
				child = n.left
			} else {
				child = n.right
			}
			return child, true
		}
		// Two children: replace with in-order successor.
		succ := minNode(n.right)
		n.key = succ.key
		n.right, _ = del(n.right, succ.key)
		removed = true
	}
	if !removed {
		return n, false
	}
	return rebalance(n), true
}

// IsBalanced reports whether the tree currently satisfies the AVL invariant.
// Useful for tests; not for normal callers.
func (t *Tree[T]) IsBalanced() bool {
	var ok bool
	ok, _ = checkBalanced(t.root)
	return ok
}

func checkBalanced[T Ordered](n *node[T]) (bool, int) {
	if n == nil {
		return true, 0
	}
	okL, hL := checkBalanced(n.left)
	if !okL {
		return false, 0
	}
	okR, hR := checkBalanced(n.right)
	if !okR {
		return false, 0
	}
	diff := hL - hR
	if diff < -1 || diff > 1 {
		return false, 0
	}
	if hL > hR {
		return true, hL + 1
	}
	return true, hR + 1
}