stl/tree.go

package stl

import (
	"math/rand/v2"
)

// Node is a node on the binary search tree.
type Node[K, V any, C Sorter[K]] struct {
	lson, rson *Node[K, V, C]
	parent     *Node[K, V, C]
	bal        uint32 // for treap balancing

	key   K
	Value V // The value on the node.
}

// Key is read-only access to the key of the node.
func (n *Node[K, V, C]) Key() K {
	return n.key
}

// Tree is a binary search tree. Always take
// addresses and don't work with values.
//
// The zero value is a usable, empty tree.
type Tree[K, V any, C Sorter[K]] struct {
	root *Node[K, V, C]
}

// TreeInc is an alias for a tree compared by the
// default Less operator.
type TreeInc[K ImplicitSortable, V any] = Tree[K, V, Less[K]]

// TreeDec is an alias for a tree compared by the
// Greater operator.
type TreeDec[K ImplicitSortable, V any] = Tree[K, V, Greater[K]]

// Clear clears the tree, erasing every element on it.
func (t *Tree[K, V, C]) Clear() {
	t.root = nil
}

// Insert tries to insert a new node on the tree, locating
// the value if the key is already in the tree.
//
// Returns the new (or found) value on the tree, and true/false
// if the key was inserted or found.
func (t *Tree[K, V, C]) Insert(key K, value V) (value_on_tree *V, is_added bool) {
	node, is_added := t.InsertNode(key, value)
	return &node.Value, is_added
}

// InsertNode does the same as Insert, but returns the entire *Node.
func (t *Tree[K, V, C]) InsertNode(key K, value V) (node *Node[K, V, C], is_added bool) {
	t.root = t.realInsertNode(t.root, nil, key, value, &node, &is_added)
	return
}

// realInsertNode recursivly searches and tries to insert the new node.
// If the key is found, it's passed into *result.
//
// Returns the new node that should be placed in CUR's place (usually itself).
func (t *Tree[K, V, C]) realInsertNode(cur, parent *Node[K, V, C], key K, value V, result **Node[K, V, C], is_added *bool) (node *Node[K, V, C]) {
	var c C
	if cur == nil {
		*result = &Node[K, V, C]{
			parent: parent,
			key:    key,
			Value:  value,
			bal:    rand.Uint32(),
		}
		*is_added = true
		return *result
	} else if c.Compare(key, cur.key) {
		// key < now.key
		cur.lson = t.realInsertNode(cur.lson, cur, key, value, result, is_added)
		return cur
	} else if c.Compare(cur.key, key) {
		// key > now.key
		cur.rson = t.realInsertNode(cur.rson, cur, key, value, result, is_added)
		return cur
	} else {
		// key == now.key
		*result = cur
		*is_added = false
		return cur
	}

}

// Find finds the value related to the key on the tree,
// or nil if it is not found.
func (t *Tree[K, V, C]) Find(key K) (value_on_tree *V) {
	node := t.FindNode(key)
	if node == nil {
		return nil
	} else {
		return &node.Value
	}
}

// FindNode finds the node instead, or nil if it is not found.
func (t *Tree[K, V, C]) FindNode(key K) (node *Node[K, V, C]) {
	return t.realFindNode(t.root, key)
}

func (t *Tree[K, V, C]) realFindNode(cur *Node[K, V, C], key K) (node *Node[K, V, C]) {
	var c C
	if cur == nil {
		return nil
	} else if c.Compare(key, cur.key) {
		// key < cur.key
		return t.realFindNode(cur.lson, key)
	} else if c.Compare(cur.key, key) {
		// key > cur.key
		return t.realFindNode(cur.rson, key)
	} else {
		// key == cur.key
		return cur
	}
}

// Deletes a node, does nothing if node is Nil.
//
// Use in tandem with FindNode to delete a key.
func (t *Tree[K, V, C]) Delete(node *Node[K, V, C]) {
	if node == nil {
		return
	}

	for node.lson != nil && node.rson != nil {
		if node.lson.bal < node.rson.bal {
			node.lson.rotate(&t.root)
		} else {
			node.rson.rotate(&t.root)
		}
	}

	if node == t.root {
		// select new root
		if node.lson != nil {
			t.root = node.lson
		} else {
			t.root = node.rson
		}
	}

	if node.lson != nil {
		node.parent.connect(node.lson, node.tell())
	} else {
		node.parent.connect(node.rson, node.tell())
	}

	// delete everything for the GC
	// Seriously, why use a tree when you have a GC?
	node.lson = nil
	node.rson = nil
	node.parent = nil
	var k K
	var v V
	node.key = k
	node.Value = v
}

func (t *Tree[K, V, C]) FirstNode() (node *Node[K, V, C]) {
	node = t.root
	if node == nil {
		return nil
	}
	for node.lson != nil {
		node = node.lson
	}
	return node
}

func (t *Tree[K, V, C]) LastNode() (node *Node[K, V, C]) {
	node = t.root
	if node == nil {
		return nil
	}
	for node.rson != nil {
		node = node.rson
	}
	return node
}

func (node *Node[K, V, C]) Next() (next *Node[K, V, C]) {
	if node.rson != nil {
		next = node.rson
		for next.lson != nil {
			next = next.lson
		}
		return
	} else {
		next = node
		for next.parent != nil && next.tell() == treeRight {
			next = next.parent
		}
		return next.parent
	}
}

func (node *Node[K, V, C]) Previous() (prev *Node[K, V, C]) {
	if node.lson != nil {
		prev = node.lson
		for prev.rson != nil {
			prev = prev.rson
		}
		return prev
	} else {
		prev = node
		for prev.parent != nil && prev.tell() == treeLeft {
			prev = prev.parent
		}
		return prev.parent
	}
}